text:
struct DataValue{T}
    hasvalue::Bool
    value::T

    DataValue{T}() where {T} = new(false)
    DataValue{T}(value::T, hasvalue::Bool=true) where {T} = new(hasvalue, value)
end

struct DataValueException <: Exception
end

const NA = DataValue{Union{}}()

DataValue(value::T, hasvalue::Bool=true) where {T} = DataValue{T}(value, hasvalue)
DataValue(value::Missing) = DataValue{Union{}}()
DataValue{T}(::Missing) where T = DataValue{T}()
DataValue() = DataValue{Union{}}()
DataValue{T}(value::Any) where T = DataValue{T}(convert(T, value), true)

Base.eltype(::Type{DataValue{T}}) where {T} = T

Base.convert(::Type{DataValue{T}}, x::DataValue{T}) where {T} = x
Base.convert(::Type{DataValue}, x::DataValue) = x

Base.convert(t::Type{DataValue{T}}, x::Any) where {T} = convert(t, convert(T, x))

function Base.convert(::Type{DataValue{T}}, x::DataValue) where {T}
    return isna(x) ? DataValue{T}() : DataValue{T}(convert(T, unsafe_get(x)))
end

Base.convert(::Type{DataValue{T}}, x::T) where {T} = DataValue{T}(x)
Base.convert(::Type{DataValue}, x::T) where {T} = DataValue{T}(x)

Base.convert(::Type{DataValue{T}}, ::Missing) where {T} = DataValue{T}()
Base.convert(::Type{DataValue}, ::Missing) = DataValue{Union{}}()

Base.convert(::Type{DataValue{T}}, ::Nothing) where {T} = DataValue{T}()
Base.convert(::Type{DataValue}, ::Nothing) = DataValue{Union{}}()

Base.convert(::Type{Union{Missing, T}}, value::DataValues.DataValue{T}) where T = get(value, missing)
Base.convert(::Type{Union{Missing, T}}, ::DataValues.DataValue{Union{}}) where T = missing
Base.convert(::Type{Any}, ::DataValue{Union{}}) = NA
Base.convert(::Type{Missing}, ::DataValue{Union{}}) = missing

Base.promote_rule(::Type{DataValue{S}}, ::Type{T}) where {S,T} = DataValue{promote_type(S, T)}
Base.promote_rule(::Type{DataValue{T}}, ::Type{Any}) where {T} = DataValue{Any}
Base.promote_rule(::Type{DataValue{Union{}}}, ::Type{Any}) = DataValue{Any}
Base.promote_rule( ::Type{Any}, ::Type{DataValue{Union{}}}) = DataValue{Any}
Base.promote_rule(::Type{DataValue{S}}, ::Type{DataValue{T}}) where {S,T} = DataValue{promote_type(S, T)}
Base.promote_op(op::Any, ::Type{DataValue{S}}, ::Type{DataValue{T}}) where {S,T} = DataValue{Base.promote_op(op, S, T)}

function Base.show(io::IO, x::DataValue{T}) where {T}
    if get(io, :compact, false)
        if isna(x)
            print(io, "#NA")
        else
            show(io, unsafe_get(x))
        end
    else
        print(io, "DataValue{")
        show(IOContext(io, :compact => true), eltype(x))
        print(io, "}(")
        if !isna(x)
            show(IOContext(io, :compact => true), unsafe_get(x))
        end
        print(io, ')')
    end
end

@inline function Base.get(x::DataValue{S}, y::T) where {S,T}
    if isbitstype(S)
        ifelse(isna(x), y, unsafe_get(x))
    else
        isna(x) ? y : unsafe_get(x)
    end
end

Base.get(x::DataValue) = isna(x) ? throw(DataValueException()) : unsafe_get(x)

"""
    getindex(x::DataValue)

Attempt to access the value of `x`. Throw a `DataValueException` if the
value is not present. Usually, this is written as `x[]`.
"""
Base.getindex(x::DataValue) = isna(x) ? throw(DataValueException()) : unsafe_get(x)

Base.get(x::DataValue{Union{}}) = throw(DataValueException())
Base.get(x::DataValue{Union{}}, y) = y

unsafe_get(x::DataValue) = x.value

isna(x) = false
isna(x::DataValue) = !x.hasvalue

hasvalue(x::DataValue) = x.hasvalue

const DataValuehash_seed = UInt === UInt64 ? 0x932e0143e51d0171 : 0xe51d0171

function Base.hash(x::DataValue, h::UInt)
    if isna(x)
        return h + DataValuehash_seed
    else
        return hash(unsafe_get(x), h)
    end
end

# TODO This is type piracy, but I think ok for now
function Base.hash(x::DataValue{Union{}}, h::UInt)
    return h + DataValuehash_seed
end

import Base.==
import Base.!=

Base.zero(::Type{DataValues.DataValue{T}}) where {T <: Number} = DataValue{T}(zero(T))
Base.zero(x::DataValues.DataValue{T}) where {T <: Number} = DataValue{T}(zero(T))
Base.zero(::Type{DataValue{T}}) where {T<:Dates.Period} = DataValue{T}(zero(T))
Base.zero(x::DataValues.DataValue{T}) where {T<:Dates.Period}= DataValue{T}(zero(T))

# C# spec section 7.10.9

==(a::DataValue{T},b::DataValue{Union{}}) where {T} = isna(a)
==(a::DataValue{Union{}},b::DataValue{T}) where {T} = isna(b)
==(a::DataValue{Union{}}, b::DataValue{Union{}}) = true
!=(a::DataValue{T},b::DataValue{Union{}}) where {T} = !isna(a)
!=(a::DataValue{Union{}},b::DataValue{T}) where {T} = !isna(b)
!=(a::DataValue{Union{}}, b::DataValue{Union{}}) = false

# Strings

for op in (:lowercase,:uppercase,:reverse,:uppercasefirst,:lowercasefirst,:chop,:chomp)
    @eval begin
        function Base.$op(x::DataValue{T}) where {T <: AbstractString}
            if isna(x)
                return DataValue{T}()
            else
                return DataValue($op(unsafe_get(x)))
            end
        end
    end
end

function Base.rstrip(f, s::DataValue{T}) where {T <: AbstractString}
    return isna(s) ? DataValue{SubString{T}}() : DataValue(rstrip(f, unsafe_get(s)))
end
function Base.rstrip(s::DataValue{T}, chars::Base.Chars) where {T <: AbstractString}
    return isna(s) ? DataValue{SubString{T}}() : DataValue(rstrip(unsafe_get(s), chars))
end
function Base.rstrip(s::DataValue{T}) where {T <: AbstractString}
    return isna(s) ? DataValue{SubString{T}}() : DataValue(rstrip(unsafe_get(s)))
end

function Base.lstrip(f, s::DataValue{T}) where {T <: AbstractString}
    return isna(s) ? DataValue{SubString{T}}() : DataValue(lstrip(f, unsafe_get(s)))
end
function Base.lstrip(s::DataValue{T}, chars::Base.Chars) where {T <: AbstractString}
    return isna(s) ? DataValue{SubString{T}}() : DataValue(lstrip(unsafe_get(s), chars))
end
function Base.lstrip(s::DataValue{T}) where {T <: AbstractString}
    return isna(s) ? DataValue{SubString{T}}() : DataValue(lstrip(unsafe_get(s)))
end

function Base.strip(s::DataValue{T}, chars::Base.Chars) where {T <: AbstractString}
    return isna(s) ? DataValue{SubString{T}}() : DataValue(strip(unsafe_get(s), chars))
end
function Base.strip(s::DataValue{T}) where {T <: AbstractString}
    return isna(s) ? DataValue{SubString{T}}() : DataValue(strip(unsafe_get(s)))
end

import Base.getindex
function Base.getindex(s::DataValue{T},i) where {T <: AbstractString}
    if isna(s)
        return DataValue{T}()
    else
        return DataValue(unsafe_get(s)[i])
    end
end

function Base.lastindex(s::DataValue{T}) where {T <: AbstractString}
    if isna(s)
        # TODO Decide whether this makes sense?
        return 0
    else
        return lastindex(unsafe_get(s))
    end
end

import Base.length
function length(s::DataValue{T}) where {T <: AbstractString}
    if isna(s)
        return DataValue{Int}()
    else
        return DataValue{Int}(length(unsafe_get(s)))
    end
end

# C# spec section 7.3.7

for op in (:+, :-, :!, :~)
    @eval begin
        import Base.$(op)
        $op(x::DataValue{T}) where {T <: Number} = isna(x) ? DataValue{T}() : DataValue($op(unsafe_get(x)))
    end
end

for op in (:+, :-, :*, :/, :%, :&, :|, :^, :<<, :>>, :min, :max)
    @eval begin
        import Base.$(op)
        $op(a::DataValue{T1},b::DataValue{T2}) where {T1 <: Number,T2 <: Number} = isna(a) || isna(b) ? DataValue{promote_type(T1,T2)}() : DataValue{promote_type(T1,T2)}($op(unsafe_get(a), unsafe_get(b)))
        $op(x::DataValue{T1},y::T2) where {T1 <: Number,T2 <: Number} = isna(x) ? DataValue{promote_type(T1,T2)}() : DataValue{promote_type(T1,T2)}($op(unsafe_get(x), y))
        $op(x::T1,y::DataValue{T2}) where {T1 <: Number,T2 <: Number} = isna(y) ? DataValue{promote_type(T1,T2)}() : DataValue{promote_type(T1,T2)}($op(x, unsafe_get(y)))
    end
end

^(x::DataValue{T},p::Integer) where {T <: Number} = isna(x) ? DataValue{T}() : DataValue(unsafe_get(x)^p)
(/)(x::DataValue{T}, y::DataValue{S}) where {T<:Integer,S<:Integer} = (isna(x) | isna(y)) ? DataValue{Float64}() : DataValue{Float64}(float(unsafe_get(x)) / float(unsafe_get(y)))
(/)(x::DataValue{T}, y::S) where {T<:Integer,S<:Integer} = isna(x) ? DataValue{Float64}() : DataValue{Float64}(float(unsafe_get(x)) / float(y))
(/)(x::T, y::DataValue{S}) where {T<:Integer,S<:Integer} = isna(y) ? DataValue{Float64}() : DataValue{Float64}(float(x) / float(unsafe_get(y)))

==(a::DataValue{T1},b::DataValue{T2}) where {T1,T2} = isna(a) && isna(b) ? true : !isna(a) && !isna(b) ? unsafe_get(a)==unsafe_get(b) : false
==(a::DataValue{T1},b::T2) where {T1,T2} = isna(a) ? false : unsafe_get(a)==b
==(a::T1,b::DataValue{T2}) where {T1,T2} = isna(b) ? false : a==unsafe_get(b)

!=(a::DataValue{T1},b::DataValue{T2}) where {T1,T2} = isna(a) && isna(b) ? false : !isna(a) && !isna(b) ? unsafe_get(a)!=unsafe_get(b) : true
!=(a::DataValue{T1},b::T2) where {T1,T2} = isna(a) ? true : unsafe_get(a)!=b
!=(a::T1,b::DataValue{T2}) where {T1,T2} = isna(b) ? true : a!=unsafe_get(b)

for op in (:<,:>,:<=,:>=)
    @eval begin
        import Base.$(op)
        $op(a::DataValue{T},b::DataValue{T}) where {T <: Number} = isna(a) || isna(b) ? false : $op(unsafe_get(a), unsafe_get(b))
        $op(x::DataValue{T1},y::T2) where {T1 <: Number,T2 <: Number} = isna(x) ? false : $op(unsafe_get(x), y)
        $op(x::T1,y::DataValue{T2}) where {T1 <: Number,T2 <: Number} = isna(y) ? false : $op(x, unsafe_get(y))
    end
end

# C# spec 7.11.4
function (&)(x::DataValue{Bool},y::DataValue{Bool})
    if isna(x)
        if isna(y) || unsafe_get(y)==true
            return DataValue{Bool}()
        else
            return DataValue(false)
        end
    elseif unsafe_get(x)==true
        return y
    else
        return DataValue(false)
    end
end

(&)(x::Bool,y::DataValue{Bool}) = x ? y : DataValue(false)
(&)(x::DataValue{Bool},y::Bool) = y ? x : DataValue(false)

function (|)(x::DataValue{Bool},y::DataValue{Bool})
    if isna(x)
        if isna(y) || !unsafe_get(y)
            return DataValue{Bool}()
        else
            return DataValue(true)
        end
    elseif unsafe_get(x)
        return DataValue(true)
    else
        return y
    end
end

(|)(x::Bool,y::DataValue{Bool}) = x ? DataValue(true) : y
(|)(x::DataValue{Bool},y::Bool) = y ? DataValue(true) : x

import Base.isless
function isless(x::DataValue{S}, y::DataValue{T}) where {S,T}
    if isna(x)
        return false
    elseif isna(y)
        return true
    else
        return isless(unsafe_get(x), unsafe_get(y))
    end
end

isless(x::S, y::DataValue{T}) where {S,T} = isna(y) ? true : isless(x, unsafe_get(y))

isless(x::DataValue{S}, y::T) where {S,T} = isna(x) ? false : isless(unsafe_get(x), y)

isless(x::DataValue{Union{}}, y::DataValue{Union{}}) = false

isless(x, y::DataValue{Union{}}) = true

isless(x::DataValue{Union{}}, y) = false

# TODO Is that the definition we want?
function Base.isnan(x::DataValue{T}) where {T<:AbstractFloat}
    return !isna(x) && isnan(unsafe_get(x))
end

# TODO Is that the definition we want?
function Base.isfinite(x::DataValue{T}) where {T<:AbstractFloat}
    return !isna(x) && isfinite(unsafe_get(x))
end

function Base.float(x::DataValue{T}) where T
    return isna(x) ? DataValue{Float64}() : DataValue{Float64}(float(unsafe_get(x)))
end

# DataValueInterfaces definitions
using DataValueInterfaces
DataValueInterfaces.nondatavaluetype(::Type{DataValue{T}}) where {T} = Union{T, Missing}
DataValueInterfaces.datavaluetype(::Type{T}) where {T <: DataValue} = T
DataValueInterfaces.datavaluetype(::Type{Union{T, Missing}}) where {T} = DataValue{T}
DataValueInterfaces.datavaluetype(::Type{Missing}) = DataValue{Union{}}


text:
export Category, FreeCategory, Ob, Hom, dom, codom, id, compose, ⋅,
  Copresheaf, FreeCopresheaf, El, ElExpr, ob, act,
  Presheaf, FreePresheaf, coact,
  DisplayedCategory, Act, hom,
  MCategory, FreeMCategory, Tight, reflexive, transitive

import Base: show

# Category
##########

""" Theory of *categories* (with no extra structure)

**Warning**: We compose functions from left to right, i.e., if f:A→B and g:B→C
then compose(f,g):A→C. Under this convention function are applied on the right,
e.g., if a∈A then af∈B.

We retain the usual meaning of the symbol ∘ (\\circ), i.e., g∘f = compose(f,g).
This usage is too entrenched to overturn, inconvenient though it may be.
We use symbol ⋅ (\\cdot) for diagrammatic composition: f⋅g = compose(f,g).
"""
@theory Category{Ob,Hom} begin
  # Unicode aliases.
  @op begin
    (→) := Hom
    (⋅) := compose
  end

  """ Object in a category """
  Ob::TYPE

  """ Morphism in a category """
  Hom(dom::Ob,codom::Ob)::TYPE

  id(A::Ob)::(A → A)
  compose(f::(A → B), g::(B → C))::(A → C) ⊣ (A::Ob, B::Ob, C::Ob)

  # Category axioms.
  ((f ⋅ g) ⋅ h == f ⋅ (g ⋅ h)
    ⊣ (A::Ob, B::Ob, C::Ob, D::Ob, f::(A → B), g::(B → C), h::(C → D)))
  f ⋅ id(B) == f ⊣ (A::Ob, B::Ob, f::(A → B))
  id(A) ⋅ f == f ⊣ (A::Ob, B::Ob, f::(A → B))
end

# Convenience constructors
compose(fs::AbstractVector) = foldl(compose, fs)
compose(f, g, h, fs...) = compose([f, g, h, fs...])

@syntax FreeCategory{ObExpr,HomExpr} Category begin
  compose(f::Hom, g::Hom) = associate_unit(new(f,g; strict=true), id)
end

function show_unicode(io::IO, expr::HomExpr{:compose}; kw...)
  Syntax.show_unicode_infix(io, expr, "⋅"; kw...)
end

function show_latex(io::IO, expr::HomExpr{:id}; kw...)
  Syntax.show_latex_script(io, expr, "\\mathrm{id}")
end
function show_latex(io::IO, expr::HomExpr{:compose}; paren::Bool=false, kw...)
  Syntax.show_latex_infix(io, expr, "\\cdot"; paren=paren)
end

function show(io::IO, ::MIME"text/plain", expr::HomExpr)
  show_unicode(io, expr)
  print(io, ": ")
  show_unicode(io, dom(expr))
  print(io, " → ")
  show_unicode(io, codom(expr))
end
function show(io::IO, ::MIME"text/latex", expr::HomExpr)
  print(io, "\$")
  show_latex(io, expr)
  print(io, " : ")
  show_latex(io, dom(expr))
  print(io, " \\to ")
  show_latex(io, codom(expr))
  print(io, "\$")
end

# (Co)presheaf
##############

""" Theory of *copresheaves*.

Axiomatized as a covariant category action.
"""
@theory Copresheaf{Ob,Hom,El} <: Category{Ob,Hom} begin
  # copresheaf = object-indexed family
  El(ob::Ob)::TYPE

  # functoriality = covariant action
  act(x::El(A), f::Hom(A,B))::El(B) ⊣ (A::Ob, B::Ob)
  @op (⋅) := act

  # action equations
  act(act(x, f), g) == act(x, (f ⋅ g)) ⊣
    (A::Ob, B::Ob, C::Ob, f::(A → B), g::(B → C), x::El(A))
  act(x, id(A)) == x ⊣ (A::Ob, x::El(A))
end

abstract type ElExpr{T} <: GATExpr{T} end

@syntax FreeCopresheaf{ObExpr,HomExpr,ElExpr} Copresheaf begin
  compose(f::Hom, g::Hom) = associate_unit(new(f,g; strict=true), id)
end

""" Theory of *presheaves*.

Axiomatized as a contravariant category action.
"""
@theory Presheaf{Ob,Hom,El} <: Category{Ob,Hom} begin
  # presheaf = object-indexed family
  El(ob::Ob)::TYPE

  # functoriality = contravariant action
  coact(f::Hom(A,B), x::El(B))::El(A) ⊣ (A::Ob, B::Ob)
  @op (⋅) := coact

  # action equations
  coact(f, coact(g, x)) == coact((f ⋅ g), x) ⊣
    (A::Ob, B::Ob, C::Ob, f::(A → B), g::(B → C), x::El(C))
  coact(id(A), x) == x ⊣ (A::Ob, x::El(A))
end

@syntax FreePresheaf{ObExpr,HomExpr,ElExpr} Presheaf begin
  compose(f::Hom, g::Hom) = associate_unit(new(f,g; strict=true), id)
end

function show(io::IO, ::MIME"text/plain", expr::ElExpr)
  show_unicode(io, expr)
  print(io, ": ")
  show_unicode(io, ob(expr))
end
function show(io::IO, ::MIME"text/latex", expr::ElExpr)
  print(io, "\$")
  show_latex(io, expr)
  print(io, " : ")
  show_latex(io, ob(expr))
  print(io, "\$")
end

function show_unicode(io::IO, expr::Union{ElExpr{:act},ElExpr{:coact}}; kw...)
  Syntax.show_unicode_infix(io, expr, "⋅"; kw...)
end
function show_latex(io::IO, expr::Union{ElExpr{:act},ElExpr{:coact}};
                    paren::Bool=false, kw...)
  Syntax.show_latex_infix(io, expr, "\\cdot"; paren=paren)
end

# Displayed category
####################

""" Theory of a *displayed category*.

More precisely, this is the theory of a base category ``C`` (`Ob`,`Hom`) and a
displayed category (`El`,`Act`) over ``C``. Displayed categories axiomatize lax
functors ``C → **Span**`` or equivalently objects in ``**Cat**/C`` in a
pleasant, generalized algebraic style.

Reference: Ahrens & Lumsdaine 2019, "Displayed categories", Definition 3.1.
"""
@theory DisplayedCategory{Ob,Hom,El,Act} <: Category{Ob,Hom} begin
  El(ob::Ob)::TYPE
  Act(hom::Hom(A,B), dom::El(A), codom::El(B))::TYPE ⊣ (A::Ob, B::Ob)

  id(x::El(A))::Act(id(A),x,x) ⊣ (A::Ob)
  compose(f̄::Act(f,x,y), ḡ::Act(g,y,z))::Act(compose(f,g),x,z) ⊣
    (A::Ob, B::Ob, C::Ob, f::(A → B), g::(B → C))

  # Displayed category axioms: category axioms over the base category.
  ((f̄ ⋅ ḡ) ⋅ h̄ == f̄ ⋅ (ḡ ⋅ h̄) ⊣
    (A::Ob, B::Ob, C::Ob, D::Ob, f::(A → B), g::(B → C), h::(C → D),
     w::El(A), x::El(B), y::El(C), z::El(D),
     f̄::Act(f,w,x), ḡ::Act(g,x,y), h̄::Act(h,y,z)))
  f̄ ⋅ id(y) == f̄ ⊣ (A::Ob, B::Ob, f::(A → B), x::El(A), y::El(B), f̄::Act(f,x,y))
  id(x) ⋅ f̄ == f̄ ⊣ (A::Ob, B::Ob, f::(A → B), x::El(A), y::El(B), f̄::Act(f,x,y))
end

# ℳ-category
############

""" Theory of an *ℳ-category*.

The term "ℳ-category", used on the
[nLab](https://ncatlab.org/nlab/show/M-category) is not very common, but the
concept itself shows up commonly. An *ℳ-category* is a category with a
distinguished wide subcategory, whose morphisms are suggestively called *tight*;
for contrast, a generic morphism is called *loose*. Equivalently, an ℳ-category
is a category enriched in the category ℳ of injections, the full subcategory of
the arrow category of Set spanned by injections.

In the following GAT, tightness is axiomatized as a property of morphisms: a
dependent family of sets over the hom-sets, each having at most one inhabitant.
"""
@theory MCategory{Ob,Hom,Tight} <: Category{Ob,Hom} begin
  Tight(hom::Hom(A,B))::TYPE ⊣ (A::Ob, B::Ob)

  # Tightness is a property.
  t == t′ ⊣ (A::Ob, B::Ob, f::Hom(A,B), t::Tight(f), t′::Tight(f))

  # Tight morphisms form a subcategory.
  reflexive(A::Ob)::Tight(id(A))
  transitive(t::Tight(f), u::Tight(g))::Tight(compose(f,g)) ⊣
    (A::Ob, B::Ob, C::Ob, f::Hom(A,B), g::Hom(B,C), t::Tight(f), u::Tight(g))
end

abstract type TightExpr{T} <: GATExpr{T} end

@syntax FreeMCategory{ObExpr,HomExpr,TightExpr} MCategory begin
  compose(f::Hom, g::Hom) = associate_unit(new(f,g; strict=true), id)
  transitive(t::Tight, u::Tight) = associate_unit(new(t,u; strict=true), reflexive)
end

ext: jl

max_forks_repo_name: epatters/CompCat


text:
using Test
using EBayes
using StructArrays
using Random

## Check single sample
ns = NormalSample(1.0)
@test isa(ns, EBayesSample)
@test isa(ns, EBayesSample{Float64})
@test response(ns) == 1.0
@test eltype(ns) == Float64
@test zero(ns) == 0
@test var(ns) == 1.0

standard_n = StandardNormalSample(1.0)
@test var(standard_n) == var(ns)
@test eltype(standard_n) == eltype(1.0)
@test zero(standard_n) == 0

ns2 = NormalSample(1.0, 1)
@test ns == ns2

# Check vectorized ebayes samples
nss = StructArray([ns; ns2])
@test isa(nss, NormalSamples)
@test isa(nss, NormalSamples{Float64})

@test isa(nss, EBayesSamples)
@test isa(nss, AbstractNormalSamples)
@test AbstractNormalSamples <: EBayesSamples
@test NormalSamples{Float64} <: EBayesSamples

nss2 = NormalSamples(nss.Z,nss.σ)
@test isa(nss2, NormalSamples)

@test response(nss) == response.(nss)
@test zeros(nss) == zeros(Float64, 2)
@test length(nss) == 2

tmp_zs = rand(10)
standard_ns = StandardNormalSample.(tmp_zs)
@test isa(standard_ns, EBayesSamples)
@test isa(standard_ns, AbstractNormalSamples)

@test response(standard_ns) == tmp_zs

## test predictions

@test predict(Normal(), 1, 0.0, ns) == 1/2
@test predict(Normal(), 1, 1.0, ns) == 1.0
@test predict(Normal(), 0, 0.0, ns) == 0.0

@test predict.(Normal(), 1, 0.0, nss) == 1/2 .* ones(2)
@test predict.(Normal(), [1;1], [0;0], nss) == 1/2 .* ones(2)

## test simulation code
Random.seed!(0)
sim1 = rand(EBayes.XieKouBrownExample1(1000))

nss_sim1 = NormalSamples(sim1)
@test isa(nss_sim1, NormalSamples)


text:
## base.jl base methods for NamedArray

## (c) 2013 David A. van Leeuwen

## This code is licensed under the MIT license
## See the file LICENSE.md in this distribution

## copy
import Base: copy, copy!
copy{T,N,AT,DT}(a::NamedArray{T,N,AT,DT}) = NamedArray{T,N,AT,DT}(copy(a.array), deepcopy(a.dicts), identity(a.dimnames))

## from array.jl
function copy!{T}(dest::NamedArray{T}, dsto::Integer, src::ArrayOrNamed{T}, so::Integer, N::
Integer)
    if so+N-1 > length(src) || dsto+N-1 > length(dest) || dsto < 1 || so < 1
        throw(BoundsError())
    end
    if isa(src, NamedArray)
        unsafe_copy!(dest.array, dsto, src.array, so, N)
    else
        unsafe_copy!(dest.array, dsto, src, so, N)
    end
end

import Base.size, Base.ndims
size(a::NamedArray) = size(a.array)
size(a::NamedArray, d) = size(a.array, d)
ndims(a::NamedArray) = ndims(a.array)

function Base.similar(a::NamedArray, t::Type, dims::Base.Dims)
    if size(a) != dims
        return NamedArray(t, dims...) # re-initialize names arrays...
    else
        return NamedArray(similar(a.array, t), a.dicts, a.dimnames)
    end
end
#similar(A::NamedArray, t::DataType, dims::Int...) = similar(A, t, dims)
#similar(A::NamedArray, t::DataType) = similar(A, t, size(A.array))
#similar(A::NamedArray) = similar(A, eltype(A.array), size(A.array))

## our own interpretation of ind2sub
Base.ind2sub(n::NamedArray, index::Integer) = tuple(map(x -> names(n, x[1])[x[2]], enumerate(ind2sub(size(n), index)))...)

## simplified text representation of namedarray
Base.writedlm(io, n::NamedVecOrMat) = writedlm(io, hcat(names(n,1), n.array))


text:
#=
    Li-ion battery modelling
 =#

 mutable struct Liion <: AbstractStorage
     # Parameters
     α_p_ch::Float64
     α_p_dch::Float64
     η_ch::Float64
     η_dch::Float64
     η_self::Float64
     α_soc_min::Float64
     α_soc_max::Float64
     lifetime::Int64
     nCycle::Float64
     bounds::NamedTuple{(:lb, :ub), Tuple{Float64, Float64}}
     # Initial conditions
     Erated_ini::Float64
     soc_ini::Float64
     soh_ini::Float64
     # Variables
     Erated::AbstractArray{Float64,2}
     carrier::Electricity
     soc::AbstractArray{Float64,3}
     soh::AbstractArray{Float64,3}
     # Eco
     cost::AbstractArray{Float64,2}
     # Inner constructor
     Liion(; α_p_ch = 1.5,
        α_p_dch = 1.5,
        η_ch = 0.9,
        η_dch = 0.9,
        η_self = 0.0005,
        α_soc_min = 0.2,
        α_soc_max = 0.8,
        lifetime = 12,
        nCycle = 2500.,
        bounds = (lb = 0., ub = 1000.),
        Erated_ini = 1e-6,
        soc_ini = 0.5,
        soh_ini = 1.) =
        new(α_p_ch, α_p_dch, η_ch, η_dch, η_self, α_soc_min, α_soc_max, lifetime, nCycle, bounds, Erated_ini, soc_ini, soh_ini)
 end

### Preallocation
 function preallocate!(liion::Liion, nh::Int64, ny::Int64, ns::Int64)
     liion.Erated = convert(SharedArray,zeros(ny+1, ns)) ; liion.Erated[1,:] .= liion.Erated_ini
     liion.carrier = Electricity()
     liion.carrier.power = convert(SharedArray,zeros(nh, ny, ns))
     liion.soc = convert(SharedArray,zeros(nh+1, ny+1, ns)) ; liion.soc[1,1,:] .= liion.soc_ini
     liion.soh = convert(SharedArray,zeros(nh+1, ny+1, ns)) ; liion.soh[1,1,:] .= liion.soh_ini
     liion.cost = convert(SharedArray,zeros(ny, ns))
     return liion
 end

 ### Operation dynamic
function compute_operation_dynamics!(h::Int64, y::Int64, s::Int64, liion::Liion, decision::Float64, Δh::Int64)
     liion.soc[h+1,y,s], liion.soh[h+1,y,s], liion.carrier.power[h,y,s] = compute_operation_dynamics(liion, (Erated = liion.Erated[y,s], soc = liion.soc[h,y,s], soh = liion.soh[h,y,s]), decision, Δh)
end

function compute_operation_dynamics(liion::Liion, state::NamedTuple{(:Erated, :soc, :soh), Tuple{Float64, Float64, Float64}}, decision::Float64, Δh::Int64)
     # Control power constraint and correction
     power_dch = max(min(decision, liion.α_p_dch * state.Erated, state.soh * state.Erated / Δh, liion.η_dch * (state.soc * (1. - liion.η_self * Δh) - liion.α_soc_min) * state.Erated / Δh), 0.)
     power_ch = min(max(decision, -liion.α_p_ch * state.Erated, -state.soh * state.Erated / Δh, (state.soc * (1. - liion.η_self * Δh) - liion.α_soc_max) * state.Erated / Δh / liion.η_ch), 0.)
     # SoC dynamic
     soc_next = state.soc * (1. - liion.η_self * Δh) - (power_ch * liion.η_ch + power_dch / liion.η_dch) * Δh / state.Erated
     # SoH dynamic
     soh_next = state.soh - (power_dch - power_ch) * Δh / (2. * liion.nCycle * (liion.α_soc_max - liion.α_soc_min) * state.Erated)
     return soc_next, soh_next, power_dch + power_ch
end

 ### Investment dynamic
 function compute_investment_dynamics!(y::Int64, s::Int64, liion::Liion, decision::Union{Float64, Int64})
     liion.Erated[y+1,s], liion.soc[1,y+1,s], liion.soh[1,y+1,s] = compute_investment_dynamics(liion, (Erated = liion.Erated[y,s], soc = liion.soc[end,y,s], soh = liion.soh[end,y,s]), decision)
 end

 function compute_investment_dynamics(liion::Liion, state::NamedTuple{(:Erated, :soc, :soh), Tuple{Float64, Float64, Float64}}, decision::Union{Float64, Int64})
     if decision > 1e-2
         Erated_next = decision
         soc_next = liion.soc_ini
         soh_next =  1.
     else
         Erated_next = state.Erated
         soc_next = state.soc
         soh_next = state.soh
     end
     return Erated_next, soc_next, soh_next
 end


text:
# Defines the abstract `Invariant` type and its interface

"""
    abstract type AbstractInvariant

An `Invariant` checks if an input satisfies some invariant.
For example, it may check whether a number is positive.

The interface of `Invariant`s is designed so that

- the invariant can be checked, given some input
- invariants can be composed to generate more complicated
    invariants
- the creation of rich error messages is possible when an
    invariant is not satisfied.

## Interface

An `Invariant` `I` must implement the following methods:

- [`title`](#)`(::I)::String`: Descriptive name for the invariant
- [`description`](#)`(::I)::String`: A longer description of the invariant,
    giving explanation and pointing to related information.
- [`satisfies`](#)`(::I, input) -> (nothing | msg)`: Check whether an `input`
    satisfies the invariant, returning either `nothing` on success
    or an error message.
"""
abstract type AbstractInvariant end

"""
    title(invariant)

Short summary of an invariant. Is used as a title in reports and error
messages
"""
function title end

"""
    description(invariant, ctx) -> Renderable

Give a more detailed description of an invariant, taking into
account the context `ctx`.
"""
description(::AbstractInvariant) = nothing
function satisfies end

function errormessage(inv::AbstractInvariant, msg)
    buf = IOBuffer()
    errormessage(IOContext(buf, :color => true, :displaysize => (88, 500)), inv, msg)
    return String(take!(buf))
end

function errormessage(io::IO, inv::AbstractInvariant, msg)
    println(io)
    showdescription(io, inv)
    println(io, msg)
end

Base.show(io::IO, inv::AbstractInvariant) = AbstractTrees.print_tree(io, inv)
AbstractTrees.printnode(io::IO, inv::AbstractInvariant) = print(io, nameof(typeof(inv)), "(\"", md(title(inv)), "\")")
AbstractTrees.children(::AbstractInvariant) = ()

# ## Basic invariant

Base.@kwdef struct Invariant <: AbstractInvariant
    fn
    title::String
    description::Union{Nothing, String} = nothing
    inputfn = identity
end

Invariant(fn, title::String; description = nothing, inputfn = identity) =
    Invariant(; fn, title, description, inputfn)

"""
    invariant(fn, title; description, inputfn)

Create an invariant with name `title` that checks an input against
`fn` which returns either `nothing` when the invariant is satisfied,
or an error message when the invariant is violated.

    invariant(inv; title, description, inputfn)

Wrap an invariant `inv`, replacing some of its attributes with the
given keyword arguments.

    invariant(invs, title, combine = :all; description)

Combine multiple invariants `invs` logically. The type of composition dependsBy default, all invariants must be
on the third argument:

- `all` (default): All invariants must be satisfied (see [`AllInvariant`](#))
- `any`: At least one invariant must be satisfied (see [`AnyInvariant`](#))
satisfied. `

## Examples

Basic usage:

{cell}
```julia
using Invariants

inv = invariant("Is negative") do n
    n < 0 ? nothing : Invariants.md("`n` is not negative!")
end
```

Successful check:

{cell}
```julia
check(inv, -1)
```

Failing check:

{cell}
```julia
check(inv, 1)
```

Or just get a Bool:

{cell}
```julia
check(Bool, inv, -1), check(Bool, inv, 1)
```

Throw an error when an invariant is not satisfied:

{cell}
```julia
check_throw(inv, 1)
```
"""
invariant(fn, title::String; kwargs...) = Invariant(fn, title; kwargs...)

invariant(title, invariants::AbstractVector{<:AbstractInvariant}; kwargs...) =
    invariant(title, invariants, all; kwargs...)
invariant(title, invariants::AbstractVector{<:AbstractInvariant},
            ::typeof(all); kwargs...) = AllInvariant(invariants, title; kwargs...)
invariant(title, invariants::AbstractVector{<:AbstractInvariant},
            ::typeof(any); kwargs...) = AnyInvariant(invariants, title; kwargs...)

title(inv::Invariant) = inv.title
description(inv::Invariant) = inv.description
satisfies(inv::Invariant, input) = inv.fn(inv.inputfn(input))

function showdescription(io, inv)
    if !isnothing(description(inv))
        println(io, description(inv))
    end
end

function showtitle(io, inv)
    print(io, md(title(inv), io))
end

function testinvariant(inv, input)
    @test_nowarn title(inv)
    @test_nowarn description(inv)
    @test_nowarn satisfies(inv, input)
end

function exampleinvariant(symbol = :n)
    return Invariant("`$symbol` is positive",
        description = "The number `$symbol` should be larger than `0`." |> md) do x
        if !(x isa Number)
            return "`$symbol` has type $(typeof(x)), but it should be a `Number` type." |> md
        else
            x > 0 && return nothing
            return "`$symbol` is not a positive number, got value `$x`. Please pass a number larger than 0." |> md
        end
    end
end

@testset "Invariant" begin
    inv = exampleinvariant()
    testinvariant(inv, 1)

    @test isnothing(satisfies(inv, 1))
    @test occursin("should be", satisfies(inv, ""))
    @test occursin("larger", satisfies(inv, -1))
    @test occursin("0", description(inv))
end


text:
module BlockMatrices

import Base:+,*,-,\,^,getindex,setindex!,show,print

# ****************************************************************
# Square Block Diagonal Matrix
# ****************************************************************

export Block, size, block_idx, broadcastf, full,
    copy, getindex, setindex!, +, -, *, \, inv, square, cache;

type Block <: AbstractMatrix{Real}

  Blocks::Array{Any}

  # This class is inefficient with too many blocks. If your block
  # diagonal matrix has many small blocks, it may be wise to
  # cache the blocks as a sparse matrix. this makes multiplication
  # an order of magnitude faster.
  mcache::SparseMatrixCSC

  Block(size::Int) = new(Array(Any,size), spzeros(0,0))
  Block(Blk::Array{Any}) = new(Blk, spzeros(0,0))
  Block(Blk::Array) = new(convert(Array{Any}, Blk), spzeros(0,0))
  Block(Blk::Array{Any}, mcache::SparseMatrixCSC) = new(Blk, mcache)

end

function Base.size(A::Block)
  if length(A.Blocks) == 0; return (0,0); end
  n = sum([size(B,1) for B in A.Blocks])
  return (n,n)
end

Base.size(A::Block, i::Integer)    = (i == 1 || i == 2) ? size(A)[1] : 1
getindex(A::Block, i::Integer)     = A.Blocks[i]
setindex!(A::Block, B, i::Integer) = begin; A.Blocks[i] = B; mcache = spzeros(0,0); end

function block_idx(A::Block)

  #block_sizes = [size(B,1) for B in A.Blocks]

  k = size(A.Blocks,1)
  IColl = Array(UnitRange, k)

  cum_count = 1
  for i = 1:k
    blk_size = size(A.Blocks[i],1);
    IColl[i] = cum_count:(cum_count + blk_size - 1)
    cum_count += blk_size
  end

  return IColl

end

function blockIter(A::Block)

  k = size(A.Blocks,1)
  cum_count = 1
  for i = 1:k
    blk_size = size(A.Blocks[i],1)
    produce(cum_count:(cum_count + blk_size - 1))
    cum_count += blk_size
  end

end

function broadcastf(op::Function, A::Block)

  B = copy(A)
  for i = 1:length(A.Blocks)
    B[i] = op(A[i])
  end
  return B

end

function broadcastf(op::Function, A::Block, B::Block)

  C = copy(A)
  for i = 1:length(A.Blocks)
    C[i] = op(A[i], B[i])
  end
  return C

end

function broadcastf(op::Function, A::Block, x::Vector)

  y = similar(x)
  i = 1
  for I = Task( () -> blockIter(A) )
    xI = view(x,I);
    y[I] = op(A.Blocks[i])
    i += 1;
  end
  return y;

end

function broadcastf(op::Function, A::Block, X::Matrix)

  Y = similar(X)
  i = 1
  for I = Task( () -> blockIter(A) )
    XI = sub(X,I,:);
    Y[I,:] = op(A.Blocks[i],XI)
    i += 1;
  end
  return Y

end

function reduce(op::Function, A::Block)

  α = 0;
  for I = Task( () -> blockIter(A) )
    α = α + op(A.Blocks[i])
  end
  return α;

end

function Base.sparse(A::Block)

  I₊, J₊, V₊ = Int[], Int[], Float64[]
  for (I,Blk) = zip(block_idx(A), A.Blocks)
    Aᵢ = sparse(Blk)
    rows = rowvals(Aᵢ)
    vals = nonzeros(Aᵢ)
    m, n = size(Aᵢ)
    for i = 1:n
       for j in nzrange(Aᵢ, i)
          row = rows[j]; val = vals[j]
          push!(J₊, i + I[1] - 1); push!(I₊, row + I[1] - 1); push!(V₊, val)
       end
    end
  end
  return sparse(I₊,J₊,V₊);

end

function cache(A::Block)
  A.mcache = sparse(A)
end

function Base.full(A::Block)

  O = zeros(size(A))
  for (I,Blk) = zip(block_idx(A), A.Blocks)
    O[I,I] = full(Blk)
  end
  return O;

end

function *(A::Block, X::Array{Float64,2})

  if length(A.mcache) != 0; return A.mcache*X; end
  Y = similar(X)
  i = 1
  for I = Task( () -> blockIter(A) )
    XI = sub(X,I,:)
    #mult!(A.Blocks[i],XI,sub(Y,I,:))
    Y[I,:] = *(A.Blocks[i],XI)
    i += 1;
  end
  return Y

end

function Base.Ac_mul_B(A::Block, X::Array{Float64,2})

  if length(A.mcache) != 0; return A.mcache*X; end
  Y = similar(X)
  i = 1
  for I = Task( () -> blockIter(A) )
    XI = sub(X,I,:)
    Y[I,:] = Ac_mul_B(A.Blocks[i],XI)
    i += 1;
  end
  return Y

end

Base.copy(A::Block)        = Block(copy(A.Blocks), copy(A.mcache))
Base.deepcopy(A::Block)    = Block(deepcopy(A.Blocks), copy(A.mcache))
+(A::Block, B::Block)      = Block(A.Blocks + B.Blocks)
-(A::Block, B::Block)      = A + (-B)
Base.inv(A::Block)         = broadcastf(inv, A)
-(A::Block)                = broadcastf(-, A)
Base.ctranspose(A::Block)  = broadcastf(ctranspose, A)
*(A::Block, B::Block)      = broadcastf(*, A, B)
Base.norm(A::Block)        = reduce(norm, A)
Base.Ac_mul_B(A::Block, B::Block) = broadcastf(Ac_mul_B, A, B)

ViewTypes = Union{SubArray}
VectorTypes = Union{Matrix, Vector, ViewTypes}

# Extra functions for dealing with views and stuff

function +(A::Diagonal, B::Block)
  i = 1
  B0 = Block(length(B.Blocks))
  for I = block_idx(B)
    dI = A.diag[I];
    B0[i] = B[i] + Diagonal(dI);
    i = i + 1
  end
  return B0
end

function mult!(A::Diagonal, x::VectorTypes, s::VectorTypes)
  for i = 1:length(x)
    s[i] = x[i]*A.diag[i];
  end
end

function mult!(A::UniformScaling, x::VectorTypes, s::VectorTypes)
  for i = 1:length(x)
    s[i] = x[i]*A.λ;
  end
end

^(A::Block,n::Integer) = broadcastf(x -> ^(x,n), A);

function show(io::IO, M::Block)
    print(io, "BLock")
end

end


text:
module TestADBM_Nmethod
    using BioEnergeticFoodWebs
    using Test

    #testing get_adbm_terms
    S = 3
    A = [0 0 0 ;
         1 1 0 ;
         0 1 1 ]
    biomass = [0.0,1.0,1.0]

    #original
    p = model_parameters(A, rewire_method = :ADBM , Z = 10.0, Nmethod = :original) #change Z to change bodymass
    ADBMterms = BioEnergeticFoodWebs.get_adbm_terms(S,p,biomass)
    #test keys
    @test collect(keys(ADBMterms)) == [:H,:λ,:E]
    #test values
    H = [Inf Inf Inf; 2.70726 Inf Inf; 2.51359 3.32226 Inf]
    @test all(isapprox.(ADBMterms[:H],H,atol = 0.001))
    λ = [0.0189 0.000259993 1.49266e-5; 0.00379282 5.21749e-5 2.99545e-6; 0.00130006 1.78839e-5 1.02674e-6]
    @test all(isapprox.(ADBMterms[:λ],λ,atol = 0.001))
    E = [1.0, 31.6228, 316.228]

    @test all(isapprox.(ADBMterms[:E],E,atol = 0.001))
    #biomass
    p_bm = model_parameters(A, rewire_method = :ADBM , Z = 10.0, Nmethod = :biomass) #change Z to change bodymass
    ADBMterms_bm = BioEnergeticFoodWebs.get_adbm_terms(S,p_bm,biomass)
    #test keys
    @test collect(keys(ADBMterms_bm)) == [:H,:λ,:E]
    H_bm = [Inf Inf Inf; 2.70726 Inf Inf; 2.51359 3.32226 Inf]
    @test all(isapprox.(ADBMterms_bm[:H],H_bm,atol = 0.001))
    λ_bm = [0.0 0.0035 0.0011 ; 0.0 0.001 0.0002 ; 0.0 0.0002 7e-5]
    @test all(isapprox.(ADBMterms_bm[:λ],λ_bm,atol = 0.001))
    E_bm = [1.0, 31.6228, 316.228]
    @test all(isapprox.(ADBMterms_bm[:E],E_bm,atol = 0.001))

end

module TestADBM_Hmethod
    using BioEnergeticFoodWebs
    using Test

    #testing get_adbm_terms
    S = 3
    A = [0 0 0 ;
         1 1 0 ;
         0 1 1 ]
    biomass = [0.0,1.0,1.0]

    #original
    p = model_parameters(A, rewire_method = :ADBM , Z = 10.0, Nmethod = :original, Hmethod = :power) #change Z to change bodymass
    ADBMterms = BioEnergeticFoodWebs.get_adbm_terms(S,p,biomass)
    #test keys
    @test collect(keys(ADBMterms)) == [:H,:λ,:E]
    #test values
    H = [1.0 31.6228 316.228 ; 31.6228 1000.0 10000.0 ; 316.228 10000.0 100000.0]
    @test all(isapprox.(ADBMterms[:H],H,atol = 0.001))
    λ = [0.0189 0.000259993 1.49266e-5; 0.00379282 5.21749e-5 2.99545e-6; 0.00130006 1.78839e-5 1.02674e-6]
    @test all(isapprox.(ADBMterms[:λ],λ,atol = 0.001))
    E = [1.0, 31.6228, 316.228]
    @test all(isapprox.(ADBMterms[:E],E,atol = 0.001))

    #biomass
    p_bm = model_parameters(A, rewire_method = :ADBM , Z = 10.0, Nmethod = :biomass, Hmethod = :power) #change Z to change bodymass
    ADBMterms_bm = BioEnergeticFoodWebs.get_adbm_terms(S,p_bm,biomass)
    #test keys
    @test collect(keys(ADBMterms_bm)) == [:H,:λ,:E]
    H_bm = [1.0 31.6228 316.228 ; 31.6228 1000.0 10000.0 ; 316.228 10000.0 100000.0]
    @test all(isapprox.(ADBMterms_bm[:H],H_bm,atol = 0.001))
    λ_bm = [0.0 0.0035 0.0011 ; 0.0 0.001 0.0002 ; 0.0 0.0002 7e-5]
    @test all(isapprox.(ADBMterms_bm[:λ],λ_bm,atol = 0.001))
    E_bm = [1.0, 31.6228, 316.228]
    @test all(isapprox.(ADBMterms_bm[:E],E_bm,atol = 0.001))

end

module TestADBM_power_ratio
    using BioEnergeticFoodWebs
    using Test

    #Testing ADBM
    S = 5
    A = [0 0 0 0 0; 0 1 0 0 0; 0 0 0 0 0; 0 1 1 0 0; 0 1 1 1 1]
    biomass = [1.0,1.0,1.0,0.0,1.0]

    #ratio method
    p = model_parameters(A, rewire_method = :ADBM, Hmethod = :ratio, Z = 10.0)
    ADBMTest = BioEnergeticFoodWebs.ADBM(S,p,biomass)
    @test ADBMTest == [0 0 0 0 0; 0 0 0 0 0; 0 0 0 0 0; 0 0 0 0 0; 1 1 1 0 0]

    #power method
    p = model_parameters(A, rewire_method = :ADBM, Hmethod = :power, Z = 10.0)
    ADBMTest = BioEnergeticFoodWebs.ADBM(S,p,biomass)
    @test ADBMTest == [0 0 0 0 0; 1 1 1 1 1; 0 0 0 0 0; 0 0 0 0 0; 1 1 1 1 1]

    #biomass Based
    p = model_parameters(A, rewire_method = :ADBM, Nmethod = :biomass, Z = 10.0)
    ADBMTest = BioEnergeticFoodWebs.ADBM(S,p,biomass)
    @test ADBMTest == [0 0 0 0 0; 0 0 0 0 0; 0 0 0 0 0; 0 0 0 0 0; 1 1 1 0 0]

end

module TestADBM_interval
    using BioEnergeticFoodWebs
    using Test
    A = [0 1 0 ; 0 0 1 ; 0 0 0]
    b = [0.8, 0.0, 0.5]
    p = model_parameters(A, rewire_method = :ADBM, adbm_trigger = :interval, adbm_interval = 100)

    s = simulate(p, b, stop = 1000)

    @test p[:extinctions] == [1,2]
    @test p[:extinctionstime] == [(0.0, 2) , (100.0, 1)]
    @test p[:tmpA] == [convert(Array{Any, 2}, A)]
    @test length(p[:rewiretime]) != 0

    A = [0 1 0 0 ; 0 0 0 1 ; 0 0 0 1 ; 0 0 0 0]
    p = model_parameters(A, rewire_method = :ADBM, adbm_trigger = :interval, adbm_interval = 1, Z = 10.0)
    b = [0.2, 0, 0.2, 1.0]

    s = simulate(p, b, stop = 10)

    @test p[:extinctions] == [2]
    @test length(p[:extinctionstime]) != 0
    @test p[:extinctionstime] == [(0.0, 2)]
    @test p[:tmpA] == map(x -> convert(Array{Any, 2}, x),[A])
    @test length(p[:rewiretime]) != 0

end


text:
using Crux
using POMDPModels
using Test
using DataStructures
using CUDA
using Flux
using Random

## data structures
@test MinHeap == MutableBinaryHeap{Float32, DataStructures.FasterForward}

## Fenwick Trees
t = FenwickTree(ones(Float32, 10))
@test t isa FenwickTree{Float32}
v = rand(16)
t2pow = FenwickTree(v)

# Inverse query on non power of 2 arrays
@test inverse_query(t, 0.5) == 1
@test inverse_query(t, 1.0) == 1
@test inverse_query(t, 0) == 1
@test inverse_query(t, -1) == 1
@test inverse_query(t, 3.0) == 3
@test inverse_query(t, 3.5) == 4
@test inverse_query(t, 4.0) == 4
@test inverse_query(t, 9.5) == 10
@test inverse_query(t, 10) == 10
@test inverse_query(t, 11) == 11

# Search up to a given range (used due to some numerical problems when searching the full range)
priorities = FenwickTree(zeros(Float32, 100000))
N = 13500
for i=1:N
    update!(priorities, i, rand(Float32))
end
ptot = prefixsum(priorities, N)

B = 320
Δp = ptot / B
inverse_query(priorities, B*Δp)
inverse_query(priorities, ptot)
@test inverse_query(priorities, B*Δp, N) <= N+1
inverse_query(priorities, ptot, N-1)

for i = 1:1000
    rns = [rand() for _=1:B]
    indices = [inverse_query(priorities, (j + rns[j] - 1) * Δp, N-1) for j=1:B]
    if !all(indices .<= N)
        i = findfirst(indices .> N)
        println("index: ", i, " rn: ", rns[i], " val: ", inverse_query(priorities, (i + rns[i] - 1) * Δp))
    end
    @test all(indices .<= N)
end

# power of 2 arrays
for i=1:1000
    rval = 5*rand()
    ind = inverse_query(t2pow, rval, 16)

    @test prefixsum(t2pow, ind) ≥ rval
    @test prefixsum(t2pow, ind-1) < rval
end
@test inverse_query(t2pow, prefixsum(t2pow, 16)) == 16

# value test
@test all(t[i] == 1 for i=1:10)
@test all(t2pow[i] ≈ v[i] for i=1:16)

# update! test
update!(t, 4, 2.5)
@test t[4] ≈ 2.5
@test prefixsum(t, 4) ≈ 5.5

## Binary heap test
v = rand(19)
b = MutableBinaryHeap{Float64, DataStructures.FasterForward}(v)
@test first(b) == minimum(v)

push!(b, 1.0)
@test b[20] == 1.

update!(b, 1, 0.)
@test first(b) == 0.00

vlarge = rand(100000)
t = FenwickTree(vlarge)
h = MutableBinaryHeap{Float64, DataStructures.FasterForward}(vlarge)

# @time first(h)
# @time push!(h, 0.5)
# @time h[999]
#
# @time t[999]
# @time inverse_query(t, 200)
# @time value(t, 999)

## mdp_data
d1 = mdp_data(ContinuousSpace(3), ContinuousSpace(4), 100)
d2 = mdp_data(ContinuousSpace(3), ContinuousSpace(4), 100, [:weight, :t, :advantage, :return, :logprob])
# @test_throws ErrorException  mdp_data(ContinuousSpace(3), ContinuousSpace(4), 100, [:bad_key]) # Now this just throws a warning
d3 = mdp_data(ContinuousSpace(3), ContinuousSpace(4), 100, ArrayType = CuArray)

@test d1[:s] == zeros(Float32, 3, 100) && d2[:s] == zeros(Float32, 3, 100) && size(d3[:s]) == (3, 100)
@test d1[:sp] == zeros(Float32, 3, 100) && d2[:sp] == zeros(Float32, 3, 100) && size(d3[:sp]) == (3, 100)
@test d1[:a] == zeros(Bool, 4, 100) && size(d2[:a]) == (4, 100) && size(d3[:a]) == (4, 100)
@test d1[:r] == zeros(Float32, 1, 100) && size(d2[:r]) == (1, 100) && size(d3[:r]) == (1, 100)
@test d1[:done] == zeros(Bool, 1, 100) && size(d2[:done]) == (1, 100) && size(d3[:done]) == (1, 100)
@test d2[:return] == zeros(1, 100) && d2[:advantage] == zeros(1, 100) && d2[:logprob] == zeros(1, 100) && d2[:t] == zeros(1, 100) && d2[:weight] == ones(1, 100)
@test !haskey(d1, :return) && !haskey(d1, :advantage) && !haskey(d1, :logprob) && !haskey(d1, :t) && !haskey(d1, :weight)
@test d3[:s] isa CuArray

## circular_indices
circ_inds(start, Nsteps, l) = mod1.(start:start + Nsteps - 1, l)
@test circ_inds(4, 60, 100) == [4:63 ...]
@test circ_inds(1, 100, 100) == [1:100 ...]
@test circ_inds(1, 101, 100) == [1:100 ..., 1]
@test circ_inds(1, 120, 100) == [1:100 ..., 1:20 ...]
@test circ_inds(90, 20, 100) == [90:100 ..., 1:9 ...]

## Get last N indices test
# A not full buffer
b = ExperienceBuffer(ContinuousSpace(2), ContinuousSpace(1), 100)
d = Dict(:s => 2*ones(2,50), :a => ones(Bool, 1,50), :sp => ones(2,50), :r => ones(1,50), :done => zeros(1,50), :weight=>zeros(1,50))
push!(b, d)

@test get_last_N_indices(b, 10) == collect(41:50)
@test get_last_N_indices(b, 1) == [50]
@test get_last_N_indices(b, 50) == collect(1:50)
@test get_last_N_indices(b, 51) == collect(1:50)
@test get_last_N_indices(b, 1000) == collect(1:50)

# full buffer
push!(b, d)
push!(b, d)
@test get_last_N_indices(b, 10) == collect(41:50)
@test get_last_N_indices(b, 1) == [50]
@test get_last_N_indices(b, 50) == collect(1:50)
@test get_last_N_indices(b, 51) == [100, collect(1:50)...]
@test get_last_N_indices(b, 100) == [collect(51:100)..., collect(1:50)...]
@test get_last_N_indices(b, 1000) == [collect(51:100)..., collect(1:50)...]

## Priority Param Construction
p = PriorityParams(100, α=0.7)
@test length(p.minsort_priorities) == 100
@test length(p.priorities) == 100
@test p.α == 0.7f0
@test p.max_priority == 1.0f0

p2 = PriorityParams(1000, p)
@test length(p2.priorities)==1000
@test p2.α == p.α
@test p2.max_priority == p.max_priority

## Construction
b = ExperienceBuffer(ContinuousSpace(2), DiscreteSpace(4), 100,)
bpriority = ExperienceBuffer(ContinuousSpace(2), DiscreteSpace(4), 50, prioritized=true)
b_gpu = b |> gpu

bcopy = ExperienceBuffer(deepcopy(b.data), elements=0)
@test b.data == bcopy.data
@test b.priority_params == bcopy.priority_params
bcopy = ExperienceBuffer(deepcopy(b.data), prioritized=true)
@test isprioritized(bcopy)

@test b isa ExperienceBuffer{Array}
@test b_gpu isa ExperienceBuffer{CuArray}

@test length(keys(b.data)) == length(keys(b_gpu.data))
@test !haskey(b, :weight)
@test :s in keys(b_gpu.data)
@test :a in keys(b_gpu.data)
@test :sp in keys(b_gpu.data)
@test :r in keys(b_gpu.data)
@test :done in keys(b_gpu.data)
@test haskey(bpriority, :weight)
@test bpriority.data[:weight] == ones(Float32, 1, 50)
@test size(b[:a]) == (4,0)

# Buffer_like function
bsmall = buffer_like(b, capacity=3, device=cpu)
@test Crux.device(bsmall) == cpu
@test keys(bsmall) == keys(b)
@test capacity(bsmall) == 3
@test length(bsmall) == 0

## Base functions
@test keys(b) == keys(b.data)

@test size(b[:s]) == (2,0)
@test length(b) == 0
@test capacity(b) == 100

@test size(b_gpu[:s]) == (2,0)
@test length(b_gpu) == 0
@test capacity(b_gpu) == 100

@test isprioritized(bpriority)
@test Crux.device(b) == cpu
@test Crux.device(bpriority) == cpu
@test Crux.device(b_gpu) == gpu

# Convert to and from GPU
@test Crux.device(gpu(b)) == gpu
@test Crux.device(cpu(b_gpu)) == cpu

## push!
#push dictionary with one element
d = Dict(:s => 2*ones(2,1), :a => ones(Bool, 4,1), :sp => ones(2,1), :r => ones(1,1), :done => zeros(1,1), :weight=>zeros(1,1))
push!(b, d)
@test length(b) == 1
@test b[:s] == 2*ones(2,1)
@test b[:a] == ones(Int, 4,1)
@test b[:sp] == ones(2,1)
@test b[:r] == ones(1,1)
@test b[:done] == zeros(1,1)

# push dictionary with more than one element
d = Dict(:s => 3*ones(2,3), :a => rand(4,3) .< 0.5, :sp => 5*ones(2,3), :r => 6*ones(1,3), :done => ones(1,3), :weight=>zeros(1,3))
push!(b, d)
@test length(b) == 4
@test b[:s][:,2:end] ==  3*ones(2,3)
@test b[:a][:,2:end] == d[:a]
@test b[:sp][:,2:end] == 5*ones(2,3)
@test b[:r][:,2:end] == 6*ones(1,3)
@test b[:done][:,2:end] == ones(1,3)

# push a buffer
push!(b, b)
@test length(b) == 8
for k in keys(b)
    @test b[k][:, 1:4] == b[k][:, 5:8]
end

## clear!
bcopy = deepcopy(b)
clear!(bcopy)
@test length(bcopy) == 0
@test bcopy.next_ind == 1

## minibatch
I = [1,2,4]
d2 = minibatch(b, I)
for k in keys(d2)
    @test all(d2[k] .== b[k][:, I])
end

## Split
@test Crux.split_batches(100, [0.5, 0.5]) == [50,50]
@test Crux.split_batches(100, [1.0]) == [100]
@test Crux.split_batches(100, [1/3, 1/3, 1/3]) == [34,33,33]
@test_throws AssertionError Crux.split_batches(100, 0.4)

v = rand(3,100)
bsplit = ExperienceBuffer(Dict(:a => v))

b1, b2 = split(bsplit, [0.5, 0.5])
@test length(b1) == 50
@test b1[:a] == v[:,1:50]
@test length(b2) == 50
@test b2[:a] == v[:,51:100]

## update_priorities!
update_priorities!(bpriority, [1,2,3], [1., 2., 3.])
@test bpriority.priority_params.max_priority == 3.0
@test bpriority.priority_params.priorities[1] ≈ 1f0^bpriority.priority_params.α
@test bpriority.priority_params.priorities[2] ≈ 2f0^bpriority.priority_params.α
@test bpriority.priority_params.priorities[3] ≈ 3f0^bpriority.priority_params.α
@test bpriority.priority_params.minsort_priorities[1] ≈ 1f0^bpriority.priority_params.α
@test bpriority.priority_params.minsort_priorities[2] ≈ 2f0^bpriority.priority_params.α
@test bpriority.priority_params.minsort_priorities[3] ≈ 3f0^bpriority.priority_params.α

push!(bpriority, d)
d
push!(bpriority, d)
@test bpriority.priority_params.max_priority == 3.0
for i=1:6
    @test bpriority.priority_params.priorities[i] ≈ 3f0^bpriority.priority_params.α
    @test bpriority.priority_params.minsort_priorities[i] ≈ 3f0^bpriority.priority_params.α
end

## sampling

# uniform sample
t = ExperienceBuffer(ContinuousSpace(2), DiscreteSpace(4), 10)
rand!(t, b)

t = ExperienceBuffer(ContinuousSpace(2), DiscreteSpace(4), 3)
Random.seed!(0)
ids = rand(1:length(b), 3)
Random.seed!(0)
rand!(t, b)

for k in keys(t)
    @test t[k] == b[k][:,ids]
end

# Test the multi-buffer sampling
t1 = ExperienceBuffer(ContinuousSpace(2), DiscreteSpace(4), 10)
d = Dict(:s => ones(2,1), :a => ones(Bool, 4,1), :sp => ones(2,1), :r => ones(1,1), :done => zeros(1,1))
push!(t1, d)

t2 = ExperienceBuffer(ContinuousSpace(2), DiscreteSpace(4), 10)
d = Dict(:s => 2*ones(2,1), :a => ones(Bool, 4,1), :sp => ones(2,1), :r => ones(1,1), :done => zeros(1,1))
push!(t2, d)

t3 = ExperienceBuffer(ContinuousSpace(2), DiscreteSpace(4), 10)
d = Dict(:s => 3*ones(2,1), :a => ones(4,1), :sp => ones(2,1), :r => ones(1,1), :done => zeros(1,1))
push!(t3, d)

t = ExperienceBuffer(ContinuousSpace(2), DiscreteSpace(4), 10)
rand!(t, t1, t2, t3)

@test all(t[:s][:, 1:4] .== 1.0)
@test all(t[:s][:, 5:7] .== 2.0)
@test all(t[:s][:, 8:10] .== 3.0)

# Priority samples
bpriority[:s] .= rand(Float32, 2, 6)
update_priorities!(bpriority, [1:6...], [1.:6. ...])
t = ExperienceBuffer(ContinuousSpace(2), DiscreteSpace(4), 1000, [:weight])
priorities = [bpriority.priority_params.priorities[i] for i=1:length(bpriority)]

rand!(t, bpriority)

# Test the the frequency of samples is proportional to their priority
length(t)
freqs = [sum(t.indices .== i) for i=1:6] ./ length(t)
probs = priorities ./ sum(priorities)
relerr = abs.(freqs .- probs) ./ probs
@test all(relerr .< 0.01)

@test all(t[:s] .== bpriority[:s][:,t.indices])
@test all(t[:weight] .<= 1.)

## merge
# bmerge = merge(b, b, capacity = 300)
# @test length(bmerge) == 200
# @test capacity(bmerge) == 300
# @test bmerge[:s][:, 1:100] == bmerge[:s][:, 101:200]

## Multi-D states
b2d = ExperienceBuffer(ContinuousSpace((2,2), UInt8), DiscreteSpace(4), 100;)
@test ndims(b2d[:s]) == 3
b2d[:s]

d = Dict(:s => 3*ones(2,2,1), :a => ones(4,1), :sp => ones(2,2,1), :r => ones(1,1), :done => zeros(1,1))
push!(b2d, d)

@test all(b2d[:s] .== 3)


text:
asvector(x::AbstractVector) = x
asvector(x) = collect(x)

"""
    pycolumntable([T=PyObject,] src) :: T

Construct a "column table" from the `Tables.jl`-compatible table `src`, namely a Python `dict` mapping column names to column vectors.
"""
function pycolumntable(::Type{T}, src) where {T}
    cols = Tables.columns(src)
    pydict(
        T,
        pystr(String(n)) => asvector(Tables.getcolumn(cols, n)) for
        n in Tables.columnnames(cols)
    )
end
pycolumntable(::Type{T}; cols...) where {T} = pycolumntable(T, cols)
pycolumntable(src) = pycolumntable(PyObject, src)
pycolumntable(; opts...) = pycolumntable(PyObject, opts)
export pycolumntable

"""
    pyrowtable([T=PyObject,] src) :: T

Construct a "row table" from the `Tables.jl`-compatible table `src`, namely a Python `list` of rows, each row being a Python `dict` mapping column names to values.
"""
function pyrowtable(::Type{T}, src) where {T}
    rows = Tables.rows(src)
    names = Tables.columnnames(rows)
    pynames = [pystr(String(n)) for n in names]
    pylist(
        T,
        pydict(pn => Tables.getcolumn(row, n) for (n, pn) in zip(names, pynames)) for
        row in rows
    )
end
pyrowtable(::Type{T}; cols...) where {T} = pyrowtable(T, cols)
pyrowtable(src) = pyrowtable(PyObject, src)
pyrowtable(; opts...) = pyrowtable(PyObject, opts)
export pyrowtable

aspandasvector(x) = asvector(x)

@init @require CategoricalArrays="324d7699-5711-5eae-9e2f-1d82baa6b597" @eval begin
    aspandasvector(x::CategoricalArrays.CategoricalArray) = begin
        codes = map(x -> x===missing ? -1 : Int(CategoricalArrays.levelcode(x))-1, x)
        cats = CategoricalArrays.levels(x)
        ordered = x.pool.ordered
        pypandasmodule().Categorical.from_codes(codes, cats, ordered=ordered)
    end
end

"""
    pypandasdataframe([T=PyObject,] [src]; ...) :: T

Construct a pandas dataframe from `src`.

Usually equivalent to `pyimport("pandas").DataFrame(src, ...)`, but `src` may also be `Tables.jl`-compatible table.
"""
pypandasdataframe(::Type{T}; opts...) where {T} = pycall(T, pypandasmodule().DataFrame; opts...)
pypandasdataframe(::Type{T}, t; opts...) where {T} = begin
    if Tables.istable(t)
        cs = Tables.columns(t)
        pycall(T, pypandasmodule().DataFrame, pydict(pystr(String(n)) => aspandasvector(Tables.getcolumn(cs, n)) for n in Tables.columnnames(cs)); opts...)
    else
        pycall(T, pypandasmodule().DataFrame, t; opts...)
    end
end
pypandasdataframe(args...; opts...) = pypandasdataframe(PyObject, args...; opts...)
export pypandasdataframe

multidict(src) = Dict(k => v for (ks, v) in src for k in (ks isa Vector ? ks : [ks]))

"""
    PyPandasDataFrame(o; indexname=:index, columntypes=(), copy=false)

Wrap the Pandas dataframe `o` as a Julia table.

This object satisfies the `Tables.jl` and `TableTraits.jl` interfaces.

- `indexname`: The name of the index column when converting this to a table, and may be `nothing` to exclude the index.
- `columntypes`: An iterable of `columnname=>type` or `[columnnames...]=>type` pairs, used when converting to a table.
- `copy`: True to copy columns on conversion.
"""
struct PyPandasDataFrame
    ref::PyRef
    indexname::Union{Symbol,Nothing}
    columntypes::Dict{Symbol,Type}
    copy::Bool
end
PyPandasDataFrame(o; indexname = :index, columntypes = (), copy = false) =
    PyPandasDataFrame(PyRef(o), indexname, multidict(columntypes), copy)
export PyPandasDataFrame

ispyreftype(::Type{PyPandasDataFrame}) = true
pyptr(df::PyPandasDataFrame) = df.ref
Base.unsafe_convert(::Type{CPyPtr}, df::PyPandasDataFrame) = checknull(pyptr(df))
C.PyObject_TryConvert__initial(o, ::Type{PyPandasDataFrame}) =
    C.putresult(PyPandasDataFrame(pyborrowedref(o)))

Base.show(io::IO, x::PyPandasDataFrame) = print(io, pystr(String, x))
Base.show(io::IO, mime::MIME, o::PyPandasDataFrame) = _py_mime_show(io, mime, o)
Base.showable(mime::MIME, o::PyPandasDataFrame) = _py_mime_showable(mime, o)

### Tables.jl / TableTraits.jl integration

Tables.istable(::Type{PyPandasDataFrame}) = true
Tables.columnaccess(::Type{PyPandasDataFrame}) = true
function Tables.columns(x::PyPandasDataFrame)
    error("not implemented")
    # # columns
    # names = Symbol[Symbol(pystr(String, c)) for c in x.o.columns]
    # columns = PyVector[PyVector{get(x.columntypes, c, missing)}(x.o[pystr(c)].to_numpy()) for c in names]
    # # index
    # if x.indexname !== nothing
    #     if x.indexname ∈ names
    #         error("table already has column called $(x.indexname), cannot use it for index")
    #     else
    #         pushfirst!(names, x.indexname)
    #         pushfirst!(columns, PyArray{get(x.columntypes, x.indexname, missing)}(x.o.index.to_numpy()))
    #     end
    # end
    # if x.copy
    #     columns = map(copy, columns)
    # end
    # return NamedTuple{Tuple(names)}(Tuple(columns))
end

IteratorInterfaceExtensions.isiterable(x::PyPandasDataFrame) = true
IteratorInterfaceExtensions.getiterator(x::PyPandasDataFrame) =
    IteratorInterfaceExtensions.getiterator(Tables.rows(x))

TableTraits.isiterabletable(x::PyPandasDataFrame) = true


text:
#=
1. get 1st value
undex0f() returns int

#=cheatsheet
nofindfirst - no problem
findfirst() extension function
    A custom remapping from the julia functions
    first()
    findfirst()
    firstIndex(v::Vector)
    firstIndex(v::Vector, i::Int)
    elementat(v::Vector, i)
    indexof(v::Vector, i)

=#
=#
#function last(index) #returns fallback / extension we're making!
#function indexin()
#firstIndex(v)
#first(v) # first : returns the contents of the first index
#Q.need for firstIndex()

#------------
#---Index

export firstIndex, indexOf, elementat

@propagate_inbounds function firstIndex(v::Vector) #
    #using first()
    return elementat(v, 1)#[1] #content
end

#checked
"""
returns index of value i
"""
@propagate_inbounds function indexOf(i, v::Vector) #where {T<:Union{Float64,Int64}}
    try
        res = findfirst(isequal(i), v)

    catch
        return -1
    end

end

#fixed
@propagate_inbounds function indexOf(i, v::Vector)
    try
        res = findfirst(isequal(i), v)
        if res isa Number
            return res
        else
            throw(error("Unexpected Error")) # 2. throw(error(ExceptionError))
        end
    catch UnexpectedError
        @error UnexpMsg exception = (UnexpectedError, catch_backtrace())
    end
    # return res
end

@propagate_inbounds function indexOf(i, v::Vector)
    try
        res = findfirst(isequal(i), v) #warning it's never used
        #typeof(res) == Nothing ? res = -1 : return Int(res) #res[1] #
    catch
        return -1
    end
    return res
end
#---    elementat ----

#checked
function elementat(v::Vector, i)
    return elementat(i, v) #findall(x -> x == i, v)[1]
end

@propagate_inbounds function elementat(i, v::Vector)
    try
        res = findall(x -> x == i, v)[1]
        length(res) == 0 ? res = -1 : res = res
    catch
        return -1
    end

    # return res
end

@propagate_inbounds function elementat(v::Vector)
    # return findall(x -> x == 1, v)[1]
    return elementat(1, v::Vector)

end

#---elements ----

#--------------------------------

#--- Indicies ----
"""
returns an array
"""
@propagate_inbounds function firstIndicies(v::Vector, i) #::AbstractArray{T,N}
    return findall(v)
    return findall(x -> x == i, v)
end

#= the rest is for the test part
#i.e. solution for these
v = [0, 1, 2, 3]
v = [1, 2, 3, 4] # alg: first(index) + last(index)

#--- testing
#benchmarking
using BenchmarkTools
v = [1, 2, 3, 4]
#a =
firstIndex(v)
=#

max_stars_repo_name: adamwillisMastery/DeepLearner

max_issues_repo_name: adamwillisMastery/DeepLearner

max_forks_repo_name: adamwillisXanax/DeepLearner


text:
mutable struct MCIntegrationVI{T<:Real} <: NumericalVI{T}
    ϵ::T #Convergence criteria
    nIter::Integer #Number of steps performed
    optimizer::LatentArray{Optimizer} #Learning rate for stochastic updates
    nMC::Int64 #Number of samples for MC Integrations
    Stochastic::Bool #Use of mini-batches
    nSamples::Int64 #Number of samples of the data
    nSamplesUsed::Int64 #Size of mini-batches
    MBIndices::AbstractVector #Indices of the minibatch
    ρ::T #Stochastic Coefficient
    HyperParametersUpdated::Bool #To know if the inverse kernel matrix must updated
    ∇η₁::AbstractVector{AbstractVector}
    ∇η₂::AbstractVector{AbstractArray}
    ν::AbstractVector{AbstractVector}
    λ::AbstractVector{AbstractVector}
    x::SubArray{T,2,Matrix{T}}#,Tuple{Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}}},true}
    y::LatentArray{SubArray}
    function MCIntegrationVI{T}(ϵ::T,nMC::Integer,nIter::Integer,optimizer::Opt,Stochastic::Bool,nSamplesUsed::Integer=1) where {T<:Real,Opt<:Optimizer}
        return new{T}(ϵ,nIter,[optimizer],[optimizer],nMC,Stochastic,1,nSamplesUsed)
    end
end

""" `MCIntegrationVI(;ϵ::T=1e-5,nMC::Integer=1000,optimizer::Optimizer=Adam(α=0.1))`

Constructor for Variational Inference via MC Integration approximation.

**Keyword arguments**

    - `ϵ::T` : convergence criteria, which can be user defined
    - `nMC::Int` : Number of samples per data point for the integral evaluation
    - `optimizer::Optimizer` : Optimizer used for the variational updates. Should be an Optimizer object from the [GradDescent.jl]() package. Default is `Adam()`
"""
function MCIntegrationVI(;ϵ::T=1e-5,nMC::Integer=1000,optimizer::Optimizer=Momentum(η=0.01)) where {T<:Real}
    MCIntegrationVI{T}(ϵ,nMC,0,optimizer,false,1)
end

""" `MCIntegrationSVI(;ϵ::T=1e-5,nMC::Integer=1000,optimizer::Optimizer=Adam(α=0.1))`

Constructor for Stochastic Variational Inference via MC integration approximation.

**Argument**

    -`nMinibatch::Integer` : Number of samples per mini-batches

**Keyword arguments**

    - `ϵ::T` : convergence criteria, which can be user defined
    - `nMC::Int` : Number of samples per data point for the integral evaluation
    - `optimizer::Optimizer` : Optimizer used for the variational updates. Should be an Optimizer object from the [GradDescent.jl]() package. Default is `Adam()`
"""
function MCIntegrationSVI(nMinibatch::Integer;ϵ::T=1e-5,nMC::Integer=200,optimizer::Optimizer=Momentum(η=0.001)) where {T<:Real}
    MCIntegrationVI{T}(ϵ,nMC,0,optimizer,true,nMinibatch)
end

function compute_grad_expectations!(model::VGP{T,L,<:MCIntegrationVI}) where {T,L}
    raw_samples = randn(model.inference.nMC,model.nLatent)
    samples = similar(raw_samples)
    for i in 1:model.nSample
        samples .= raw_samples.*[sqrt(model.Σ[k][i,i]) for k in 1:model.nLatent]' .+ [model.μ[k][i] for k in 1:model.nLatent]'
        grad_samples(model,samples,i)
    end
end

function compute_grad_expectations!(model::SVGP{T,L,<:MCIntegrationVI}) where {T,L}
    raw_samples = randn(model.inference.nMC,model.nLatent)
    samples = similar(raw_samples)
    Σ = opt_diag.(model.κ.*model.Σ,model.κ)
    μ = model.κ.*model.μ
    for i in 1:model.nSample
        samples .= raw_samples.*[sqrt(Σ[k][i]) for k in 1:model.nLatent]' .+ [μ[k][i] for k in 1:model.nLatent]'
        grad_samples(model,samples,i)
    end
end

function compute_log_expectations(model::VGP{T,L,<:MCIntegrationVI}) where {T,L}
    raw_samples = randn(model.inference.nMC,model.nLatent)
    samples = similar(raw_samples)
    loglike = 0.0
    for i in 1:model.nSample
        samples .= raw_samples.*[sqrt(model.Σ[k][i,i]) for k in 1:model.nLatent]' .+ [model.μ[k][i] for k in 1:model.nLatent]'
        loglike += mean(mapslices(f->log(pdf(model.likelihood,model.likelihood.y_class[i],f)),samples,dims=2))
    end
    return loglike
end

function compute_log_expectations(model::SVGP{T,L,<:MCIntegrationVI}) where {T,L}
    raw_samples = randn(model.inference.nMC,model.nLatent)
    samples = similar(raw_samples)
    loglike = 0.0
    Σ = opt_diag.(model.κ.*model.Σ,model.κ)
    μ = model.κ.*model.μ
    for i in model.inference.MBIndices
        samples .= raw_samples.*[sqrt(Σ[k][i]) for k in 1:model.nLatent]' .+ [μ[k][i] for k in 1:model.nLatent]'
        loglike += mean(mapslices(f->log(pdf(model.likelihood,model.likelihood.y_class[i],f)),samples,dims=2))
    end
    return model.inference.ρ*loglike
end


text:
#================================================================

Common matrix and vector operations used in multiple locations
throughout the module.

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

# Compute the norms of vectors and columns of matrices
col_norms(x::Union{AbstractVector,AbstractMatrix}) =
    map(norm, eachcol(x))

col_norms(x::Union{Vector,Matrix}) =
    map(BLAS.nrm2, eachcol(x))

col_norms(x::SparseVector) =
    [BLAS.nrm2(x.nzval)]

col_norms(x::SparseMatrixCSC{T}) where {T} = begin
    output = Vector{T}(undef, size(x,2))
    @inbounds for ii = 1:size(x,2)
        result = T(0)
        start_idx, end_idx = x.colptr[ii], x.colptr[ii+1]-1
        @simd for idx = start_idx:end_idx
            result += x.nzval[idx].^2
        end
        output[ii] = √result
    end
    return output
end


text:
@testset "splitdims" begin
    # Matrix
    @test splitdims([1 2; 3 4]) == [[1,3], [2,4]]
    @test splitdims([1 2; 3 4], 1) == [[1,2], [3,4]]
    @test splitdims([1 2; 3 4], 2) == [[1,3], [2,4]]
    @test splitdims([1 2; 3 4], (1,)) == [[1,2], [3,4]]
    @test splitdims([1 2; 3 4], (2,)) == [[1,3], [2,4]]
    @test splitdims([1 2; 3 4], ()) == fill([1 2; 3 4], ())
    tmp = Matrix{Array{Int, 0}}(undef,2,2) # Stupid hvcat removes the Array{0}'s...
    tmp[1,1] = fill(1, ()); tmp[1,2] = fill(2, ()); tmp[2,1] = fill(3, ()); tmp[2,2] = fill(4, ())
    @test splitdims([1 2; 3 4], (1, 2)) == tmp
    tmp[1,2][] = 3; tmp[2,1][] = 2
    @test splitdims([1 2; 3 4], (2, 1)) == tmp

    # Vector
    @test splitdims([1,2,3]) == [fill(1, ()), fill(2, ()), fill(3, ())]
    @test splitdims([1,2,3], (1,)) == [fill(1, ()), fill(2, ()), fill(3, ())]
    @test splitdims([1,2,3], ()) == fill([1,2,3], ())

    # Array{0}
    @test splitdims(fill(1, ())) == fill(fill(1, ()), ())
end

@testset "splitdimsview" begin
    # Matrix
    @test splitdimsview([1 2; 3 4]) == [[1,3], [2,4]]
    @test splitdimsview([1 2; 3 4], 1) == [[1,2], [3,4]]
    @test splitdimsview([1 2; 3 4], 2) == [[1,3], [2,4]]
    @test splitdimsview([1 2; 3 4], (1,)) == [[1,2], [3,4]]
    @test splitdimsview([1 2; 3 4], (2,)) == [[1,3], [2,4]]
    @test splitdimsview([1 2; 3 4], ()) == fill([1 2; 3 4], ())
    tmp = Matrix{Array{Int, 0}}(undef,2,2) # Stupid hvcat removes the Array{0}'s...
    tmp[1,1] = fill(1, ()); tmp[1,2] = fill(2, ()); tmp[2,1] = fill(3, ()); tmp[2,2] = fill(4, ())
    @test splitdimsview([1 2; 3 4], (1, 2)) == tmp
    tmp[1,2][] = 3; tmp[2,1][] = 2
    @test splitdimsview([1 2; 3 4], (2, 1)) == tmp

    # Vector
    @test splitdimsview([1,2,3]) == [fill(1, ()), fill(2, ()), fill(3, ())]
    @test splitdimsview([1,2,3], (1,)) == [fill(1, ()), fill(2, ()), fill(3, ())]
    @test splitdimsview([1,2,3], ()) == fill([1,2,3], ())

    # Array{0}
    @test splitdimsview(fill(1, ())) == fill(fill(1, ()), ())
end


text:
# syntax: proto3
using ProtoBuf
import ProtoBuf.meta

mutable struct Imputer <: ProtoType
    __protobuf_jl_internal_meta::ProtoMeta
    __protobuf_jl_internal_values::Dict{Symbol,Any}
    __protobuf_jl_internal_defaultset::Set{Symbol}

    function Imputer(; kwargs...)
        obj = new(meta(Imputer), Dict{Symbol,Any}(), Set{Symbol}())
        values = obj.__protobuf_jl_internal_values
        symdict = obj.__protobuf_jl_internal_meta.symdict
        for nv in kwargs
            fldname, fldval = nv
            fldtype = symdict[fldname].jtyp
            (fldname in keys(symdict)) || error(string(typeof(obj), " has no field with name ", fldname))
            if fldval !== nothing
                values[fldname] = isa(fldval, fldtype) ? fldval : convert(fldtype, fldval)
            end
        end
        obj
    end
end # mutable struct Imputer
const __meta_Imputer = Ref{ProtoMeta}()
function meta(::Type{Imputer})
    ProtoBuf.metalock() do
        if !isassigned(__meta_Imputer)
            __meta_Imputer[] = target = ProtoMeta(Imputer)
            fnum = Int[1,2,3,4,5,6,7,11,12,13]
            allflds = Pair{Symbol,Union{Type,String}}[:imputedDoubleValue => Float64, :imputedInt64Value => Int64, :imputedStringValue => AbstractString, :imputedDoubleArray => DoubleVector, :imputedInt64Array => Int64Vector, :imputedStringDictionary => StringToDoubleMap, :imputedInt64Dictionary => Int64ToDoubleMap, :replaceDoubleValue => Float64, :replaceInt64Value => Int64, :replaceStringValue => AbstractString]
            oneofs = Int[1,1,1,1,1,1,1,2,2,2]
            oneof_names = Symbol[Symbol("ImputedValue"),Symbol("ReplaceValue")]
            meta(target, Imputer, allflds, ProtoBuf.DEF_REQ, fnum, ProtoBuf.DEF_VAL, ProtoBuf.DEF_PACK, ProtoBuf.DEF_WTYPES, oneofs, oneof_names)
        end
        __meta_Imputer[]
    end
end
function Base.getproperty(obj::Imputer, name::Symbol)
    if name === :imputedDoubleValue
        return (obj.__protobuf_jl_internal_values[name])::Float64
    elseif name === :imputedInt64Value
        return (obj.__protobuf_jl_internal_values[name])::Int64
    elseif name === :imputedStringValue
        return (obj.__protobuf_jl_internal_values[name])::AbstractString
    elseif name === :imputedDoubleArray
        return (obj.__protobuf_jl_internal_values[name])::DoubleVector
    elseif name === :imputedInt64Array
        return (obj.__protobuf_jl_internal_values[name])::Int64Vector
    elseif name === :imputedStringDictionary
        return (obj.__protobuf_jl_internal_values[name])::StringToDoubleMap
    elseif name === :imputedInt64Dictionary
        return (obj.__protobuf_jl_internal_values[name])::Int64ToDoubleMap
    elseif name === :replaceDoubleValue
        return (obj.__protobuf_jl_internal_values[name])::Float64
    elseif name === :replaceInt64Value
        return (obj.__protobuf_jl_internal_values[name])::Int64
    elseif name === :replaceStringValue
        return (obj.__protobuf_jl_internal_values[name])::AbstractString
    else
        getfield(obj, name)
    end
end

export Imputer


text:
# In this example, we use EKI to tune the closure parameters of a HydrostaticFreeSurfaceModel
# with a CATKEVerticalDiffusivity closure in order to align the predictions of the model
# to those of a high-resolution LES data generated in LESbrary.jl. Here `predictions` refers to the
# 1-D profiles of temperature, velocity, and turbulent kinetic energy horizontally averaged over a
# 3-D physical domain.

pushfirst!(LOAD_PATH, joinpath(@__DIR__, "../.."))

using Oceananigans
using LinearAlgebra, Distributions, JLD2, DataDeps, Random
using Oceananigans.Units
using Oceananigans.TurbulenceClosures: CATKEVerticalDiffusivity, RiBasedVerticalDiffusivity
using OceanBoundaryLayerParameterEstimation
using ParameterEstimocean
using ParameterEstimocean.Parameters: closure_with_parameters
using ParameterEstimocean.PseudoSteppingSchemes
using ParameterEstimocean.EnsembleKalmanInversions: eki_objective

Random.seed!(1234)

Nz = 32
Nensemble = 128
architecture = GPU()
# prior_type = "scaled_logit_normal"
prior_type = "normal"

path = joinpath(directory, "results.txt")
o = open_output_file(path)
write(o, "Training relative weights: $(calibration.relative_weights) \n")
write(o, "Validation relative weights: $(validation.relative_weights) \n")
write(o, "Training default parameters: $(validation.default_parameters) \n")
write(o, "Validation default parameters: $(validation.default_parameters) \n")

write(o, "------------ \n \n")
default_parameters = calibration.default_parameters
train_loss_default = calibration(default_parameters)
valid_loss_default = validation(default_parameters)
write(o, "Default parameters: $(default_parameters) \nLoss on training: $(train_loss_default) \nLoss on validation: $(valid_loss_default) \n------------ \n \n")

#####
##### Set up ensemble model
#####

begin
    Δt = 5minutes
    field_names = (:b, :u, :v, :e)
    fields_by_case = Dict(
    "free_convection" => (:b, :e),
    "weak_wind_strong_cooling" => (:b, :u, :v, :e),
    "strong_wind_weak_cooling" => (:b, :u, :v, :e),
    "strong_wind" => (:b, :u, :v, :e),
    "strong_wind_no_rotation" => (:b, :u, :e)
    )

    parameter_set = CATKEParametersRiDependent

    parameter_names = (:CᵂwΔ,  :Cᵂu★, :Cᴰ,
                       :Cˢc,   :Cˢu,  :Cˢe,
                       :Cᵇc,   :Cᵇu,  :Cᵇe,
                       :Cᴷc⁻,  :Cᴷu⁻, :Cᴷe⁻,
                       :Cᴷcʳ,  :Cᴷuʳ, :Cᴷeʳ,
                       :CᴷRiᶜ, :CᴷRiʷ)

    parameter_set = ParameterSet{CATKEVerticalDiffusivity}(Set(parameter_names),
                                nullify = Set([:Cᴬu, :Cᴬc, :Cᴬe]))

    transformation = (b = Transformation(normalization=ZScore()),
                      u = Transformation(normalization=ZScore()),
                      v = Transformation(normalization=ZScore()),
                      e = Transformation(normalization=RescaledZScore(1e-1)))

    closure = closure_with_parameters(CATKEVerticalDiffusivity(Float64;), parameter_set.settings)

    directory = "calibrate_catke_to_lesbrary/"
    isdir(directory) || mkpath(directory)
end

#####
##### Build free parameters
#####

function build_prior(name)
    b = bounds(name, parameter_set)
    prior_type == "scaled_logit_normal" && return ScaledLogitNormal(bounds=b)
    prior_type == "normal" && return Normal(mean(b), -(b...)/6)
end

free_parameters = FreeParameters(named_tuple_map(parameter_set.names, build_prior))

#####
##### Build the Inverse Problem
#####

output_map = ConcatenatedOutputMap()

function inverse_problem(Nensemble, times)
    observations = SyntheticObservationsBatch(six_day_suite_path_2m, times, Nz; architecture, transformation, field_names, fields_by_case)
    simulation = lesbrary_ensemble_simulation(observations; Nensemble, architecture, closure, Δt)
    ip = InverseProblem(observations, simulation, free_parameters; output_map)
    return ip
end

training_times = [1.0day, 1.5days, 2.0days, 2.5days, 3.0days]
validation_times = [3.0days, 3.5days, 4.0days]
testing_times = [4.0days, 4.5days, 5.0days, 5.5days, 6.0days]

training = inverse_problem(Nensemble, training_times)
validation = inverse_problem(Nensemble, validation_times)
testing = inverse_problem(Nensemble, testing_times)

y = observation_map(training);
θ = named_tuple_map(parameter_set.names, name -> default(name, parameter_set))
G = forward_map(training, [θ])
zc = [mapslices(norm, G .- y, dims = 1)...]

###
### Calibrate
###

iterations = 3

function estimate_noise_covariance(times)
    observation_high_res = SyntheticObservationsBatch(six_day_suite_path_1m, times, Nz; architecture, transformation, field_names, fields_by_case)
    observation_mid_res = SyntheticObservationsBatch(six_day_suite_path_2m, times, Nz; architecture, transformation, field_names, fields_by_case)
    observation_low_res = SyntheticObservationsBatch(six_day_suite_path_4m, times, Nz; architecture, transformation, field_names, fields_by_case)

    Nobs = Nz * (length(times) - 1) * sum([length(obs.forward_map_names) for obs in observation_low_res])
    noise_covariance = estimate_η_covariance(output_map, [observation_low_res, observation_mid_res, observation_high_res]) .+ Matrix(1e-10 * I, Nobs, Nobs)
    return noise_covariance
end

# noise_covariance = estimate_noise_covariance(training_times)

noise_covariance = 1e-2

resampler = Resampler(acceptable_failure_fraction=0.5, only_failed_particles=true)

pseudo_stepping = Iglesias2021()
eki = EnsembleKalmanInversion(training; noise_covariance, pseudo_stepping, resampler, tikhonov = true)
iterate!(eki; iterations = 10, show_progress=false, pseudo_stepping)
visualize!(training, eki.iteration_summaries[end].ensemble_mean;
    field_names = [:u, :v, :b, :e],
    directory = pwd(),
    filename = "realizations_training_estimated_gamma_y.png"
)
visualize!(validation, eki.iteration_summaries[end].ensemble_mean;
    field_names = [:u, :v, :b, :e],
    directory = pwd(),
    filename = "realizations_validation_estimated_gamma_y.png"
)
visualize!(testing, eki.iteration_summaries[end].ensemble_mean;
    field_names = [:u, :v, :b, :e],
    directory = pwd(),
    filename = "realizations_testing_estimated_gamma_y.png"
)

plot_parameter_convergence!(eki, pwd())
plot_error_convergence!(eki, pwd())

# eki = EnsembleKalmanInversion(training; noise_covariance=1e-2, pseudo_stepping, resampler, tikhonov = true)
# iterate!(eki; iterations = 10, show_progress=false, pseudo_stepping)
# visualize!(training, eki.iteration_summaries[end].ensemble_mean;
#     field_names = [:u, :v, :b, :e],
#     directory = pwd(),
#     filename = "realizations_training_constant_gamma_y.png"
# )
# visualize!(validation, eki.iteration_summaries[end].ensemble_mean;
#     field_names = [:u, :v, :b, :e],
#     directory = pwd(),
#     filename = "realizations_validation_constant_gamma_y.png"
# )

# for (pseudo_scheme, name) in zip([Default(cov_threshold=0.01), ConstantConvergence(convergence_ratio=0.7), Kovachki2018InitialConvergenceThreshold(), Iglesias2021(), GPLineSearch()],
#                                 ["default", "constant_conv", "kovachki_2018", "iglesias2021", "gp_linesearch"])

#     @show name
#     eki = EnsembleKalmanInversion(training; noise_covariance, resampler, tikhonov = true)
#     iterate!(eki; iterations = 10, show_progress=false, pseudo_stepping = pseudo_scheme)

#     dir = directory * "_" * name
#     visualize!(training, eki.iteration_summaries[end].ensemble_mean;
#         field_names = [:u, :v, :b, :e],
#         directory = dir,
#         filename = "realizations_training.png"
#     )
#     # visualize!(validation, eki.iteration_summaries[end].ensemble_mean;
#     #     field_names = [:u, :v, :b, :e],
#     #     directory = dir,
#     #     filename = "realizations_validation.png"
#     # )
#     # visualize!(testing, eki.iteration_summaries[end].ensemble_mean;
#     #     field_names = [:u, :v, :b, :e],
#     #     directory = dir,
#     #     filename = "realizations_testing.png"
#     # )

#     plot_parameter_convergence!(eki, dir)
#     plot_pairwise_ensembles!(eki, dir)
#     plot_error_convergence!(eki, dir)
# end

# # ###
# # ### Summary Plots
# # ###

# # plot_parameter_convergence!(eki, directory)
# # plot_pairwise_ensembles!(eki, directory)
# # plot_error_convergence!(eki, directory)

# # visualize!(training, eki.iteration_summaries[end].ensemble_mean;
# #     field_names = [:u, :v, :b, :e],
# #     directory,
# #     filename = "realizations.png"
# # )

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

# # validation_noise_covariance = estimate_noise_covariance(validation_times)
# # function validation_loss_final(pseudo_stepping)
# #     eki = EnsembleKalmanInversion(training; noise_covariance, pseudo_stepping, resampler)
# #     θ_end = iterate!(eki; iterations, pseudo_stepping)
# #     θ_end = collect(θ_end)

# #     eki_validation = EnsembleKalmanInversion(validation; noise_covariance = validation_noise_covariance, pseudo_stepping, resampler)
# #     G_end_validation = forward_map(validation, θ_end)[:, 1]

# #     # Vector of (Φ₁, Φ₂) pairs, one for each ensemble member at the current iteration
# #     # objective_values = [eki_objective(eki_validation, θ[j], G[:, j]; inv_sqrt_Γθ, constrained=true) for j in 1:size(G, 2)]
# #     # validation_loss_per_iteration = sum.(objective_values)

# #     loss_final = sum(eki_objective(eki_validation, θ_end, G_end_validation;
# #                                                 constrained=true))

# #     return loss_final
# # end

# # # function testing_loss_trajectory(pseudo_stepping)
# # #     eki_testing = EnsembleKalmanInversion(testing; noise_covariance, pseudo_stepping, resampler)
# # #     G_end_testing = forward_map(testing, θ_end)

# # #     # Run EKI to train

# # #     # Vector of (Φ₁, Φ₂) pairs, one for each ensemble member at the current iteration
# # #     objective_values = [eki_objective(eki_testing, θ[j], G[:, j]; inv_sqrt_Γθ, constrained=true) for j in 1:size(G, 2)]
# # #     testing_loss_per_iteration = sum.(objective_values)
# # # end

# # optim_iterations = 10

# # using Optim
# # using Optim: minimizer

# # frobenius_norm(A) = sqrt(sum(A .^ 2))

# # using ParameterEstimocean.PseudoSteppingSchemes: observations, obs_noise_covariance, inv_obs_noise_covariance
# # function kovachki_2018_update2(Xₙ, Gₙ, eki; Δtₙ=1.0)

# #     y = observations(eki)
# #     Γy = obs_noise_covariance(eki)

# #     N_ens = size(Xₙ, 2)
# #     g̅ = mean(G, dims = 2)
# #     Γy⁻¹ = inv_obs_noise_covariance(eki)

# #     # Fill transformation matrix (D(uₙ))ᵢⱼ = ⟨ G(u⁽ⁱ⁾) - g̅, Γy⁻¹(G(u⁽ʲ⁾) - y) ⟩
# #     D = zeros(N_ens, N_ens)
# #     for j = 1:N_ens, i = 1:N_ens
# #         D[i, j] = dot(Gₙ[:, j] - g̅, Γy⁻¹ * (Gₙ[:, i] - y))
# #     end

# #     # Update
# #     Xₙ₊₁ = Xₙ - Δtₙ * Xₙ * D

# #     return Xₙ₊₁
# # end

# # ##
# # ## Make sure kovachki_2018 agrees with iglesias_2013
# # ##

# # # using ParameterEstimocean.PseudoSteppingSchemes: iglesias_2013_update, kovachki_2018_update
# # # Gⁿ = eki.forward_map_output
# # # Xⁿ = eki.unconstrained_parameters

# # # r = iglesias_2013_update(Xⁿ, Gⁿ, eki; Δtₙ=1.0)
# # # t = kovachki_2018_update2(Xⁿ, Gⁿ, eki; Δtₙ=1.0)

# # # f(x) = (x-0.5)^2
# # # result2 = optimize(f, 1e-10, 1.0, Brent(); iterations=optim_iterations, store_trace=true, extended_trace=true)

# # f(step_size) = validation_loss_final(Constant(; step_size))
# # # result = optimize(f, 1e-10, 1.0, Brent(); iterations=optim_iterations, store_trace=true)
# # # p = minimizer(result)

# # f_log(step_size) = validation_loss_final(Constant(; step_size = 10^(step_size)))
# # # result = optimize(f_log, -3, 0.0, Brent(); iterations=optim_iterations, store_trace=true)
# # # p = 10^(minimizer(result))
# # # @show Optim.x_trace(result)
# # # @show 10 .^ (Optim.x_trace(result))
# # # @show Optim.f_trace(result)

# # # a = [f_log(step_size) for step_size = -3.0:0.5:0.0]
# # # b = [f(step_size) for step_size = 0.1:0.1:1.0]

# # # using CairoMakie
# # # fig = Figure()
# # # lines(fig[1,1], collect(-3.0:0.5:0.0), a)
# # # lines(fig[1,2], collect(0.1:0.1:1.0), b)
# # # save(joinpath(directory, "1d_loss_landscape.png"), fig)

# # # f(convergence_ratio) = validation_loss_final(ConstantConvergence(; convergence_ratio))
# # # result = optimize(f, 0.1, 1.0, Brent(); iterations=optim_iterations, store_trace=true)
# # # p = minimizer(result)

# # # f(initial_step_size) = validation_loss_final(Kovachki2018(; initial_step_size))
# # # result = optimize(f, 1e-10, 1.0, Brent(); iterations=optim_iterations, store_trace=true)
# # # p = minimizer(result)

# # # f(cov_threshold) = validation_loss_final(Default(; cov_threshold = 10^(cov_threshold)))
# # # result = optimize(f, -3, 0.0, Brent(); iterations=optim_iterations, store_trace=true)
# # # p = 10 .^ (minimizer(result))

# # # f(learning_rate) = validation_loss_final(GPLineSearch(; learning_rate = 10^(learning_rate)))
# # # result = optimize(f, -3, 0.0, Brent(); iterations=optim_iterations, store_trace=true)
# # # p = 10 .^ (minimizer(result))

# # # pseudo_stepping = Constant(; step_size=1.0)
# # # # using StatProfilerHTML
# # # # @profilehtml parameters = iterate!(eki; iterations)
# # # @time parameters = iterate!(eki; iterations, pseudo_stepping)
# # # visualize!(training, parameters;
# # #     field_names = [:u, :v, :b, :e],
# # #     directory,
# # #     filename = "perfect_model_visual_calibrated.png"
# # # )
# # # @show parameters

include("emulate_sample.jl")


text:
using DynamicGrids, Test
using DynamicGrids: inbounds

@testset "boundary overflow checks are working" begin
    @testset "inbounds with RemoveOverflow() returns index and false for an overflowed index" begin
        @test inbounds((1, 1), (4, 5), RemoveOverflow()) == ((1,1),true)
        @test inbounds((2, 3), (4, 5), RemoveOverflow()) == ((2,3),true)
        @test inbounds((4, 5), (4, 5), RemoveOverflow()) == ((4,5),true)
        @test inbounds((-3, -100), (4, 5), RemoveOverflow()) == ((-3,-100),false)
        @test inbounds((0, 0), (4, 5), RemoveOverflow()) == ((0,0),false)
        @test inbounds((2, 3), (3, 2), RemoveOverflow()) == ((2,3),false)
        @test inbounds((2, 3), (1, 4), RemoveOverflow()) == ((2,3),false)
        @test inbounds((200, 300), (2, 3), RemoveOverflow()) == ((200,300),false)
    end
    @testset "inbounds with WrapOverflow() returns new index and true for an overflowed index" begin
        @test inbounds((-2,3), (10, 10), WrapOverflow()) == ((8,3),true)
        @test inbounds((2,0), (10, 10), WrapOverflow()) == ((2,10),true)
        @test inbounds((22,0), (10, 10), WrapOverflow()) == ((2,10),true)
        @test inbounds((-22,0), (10, 10), WrapOverflow()) == ((8,10),true)
    end
end


text:
# ## Environmental sensors with the SenseHat
# ### Reads and displays temperature, humidity and pressure.

using SenseHat

# Create channel to store stick events
c = Channel{StickEvent}(32)
@async while true
   # Push every event into the channel
   put!(c, readstick())
end

while true
   # Retrieve events from channel
    ev = take!(c)
    if ev.direc == SenseHat.Stick.MIDDLE && ev.state == SenseHat.Stick.PRESS
      # Flash the Julia logo on the display
      LED[:] = SenseHat.JULIA_LOGO; sleep(1); led_clear()
      # Display atmospheric pressure
      show_message(string("P R E S S U R E:", round(Int, pressure()), "mb"), 0.08, colorant"purple")
      # Display temperature
      show_message(string("T E M P E R A T U R E:", round(Int, temperature()), "C"), 0.08, colorant"red")
      # Display humidity
      show_message(string("H U M I D I T Y:", round(Int, humidity()), "%"), 0.08, colorant"blue")
    end
end


text:
using RetroSignalModel
using ModelingToolkit
using DataFrames
using Optim
using CSV

versioninfo()

optim_params(targetratio=5)

res = optim_params_threads(120; targetratio=5);

paramaps = map(res) do x
    Dict(Symbol(k) => v for (k, v) in x.parammap)
end

CSV.write("solution_rtgMTK_optim_5x.csv", DataFrame(paramaps))


text:
# For, While, Do, If, CompoundExpression

# FIXME.  We should use dynamically scoped local variables for Do.
# Block is automatically used to localize values of iterators in iteration constructs such as Do, Sum, and Table
# Hmmm. looks like I noted this a long time ago in the doc for Do.

# FIXME. We need to delete local variables before breaking or returning.
macro checkbreak(arg)
    return esc(:( begin
      if is_throw()
        return($arg)
    end
    if is_break()
        clear_break()
        break
    end end))
end

macro unsetbreak()
   esc(:(clear_break()))
end

# Get Return(x), we return Return(x), rather than x.
macro checkthrowreturn(mx)
    esc(:(is_Mxpr($mx, :Return) && return $mx))
end

# Get Return(x), we return x, or Null
macro checkreturn(mx)
    esc(:(is_Mxpr($mx, :Return) && return length($mx) == 0 ? Null : $mx[1]))
end

macro checkthrowcontinue(mx)
    esc(:(is_Mxpr($mx, :Continue) && return $mx))
end

# Probably need to check for
macro checkcontinue(mx, incr)
    esc(:( if is_Mxpr($mx, :Continue) begin
           doeval(incr)
           continue
           end
           end
           ))
end

macro checkcontinue0(mx)
    esc( :( if is_Mxpr($mx, :Continue) continue end ))
end

# NOTE. See comment at top. We should use dynamcic scoping here.
# Localize variables.
# For lexically scoped variables. Replace symbol os with ns in ex
# We should follow what we did in table and set the value in the function,
# rather than localizing the variable and letting meval do the setting.
function substitute_symbol(ex, os, ns)
    if is_Mxpr(ex)
        args = margs(ex)
        @inbounds for i in 1:length(args)
            args[i] = substitute_symbol(args[i], os, ns)
        end
    end
    if ex == os
        return ns
    else
        return ex
    end
end

### For

@mkapprule For :nargs => 3:4

@sjdoc For """
    For(start, test, incr, body)

performs a "for" loop.

For example, `For(i=1, i<=3,  Increment(i) , Println(i))`.

The variable `i` is not local to the `For` loop.

Using `Increment(i)` is currently much faster than `i = i + 1` (but what about `i += 1` ?)
"""

# This is pretty fast: For(i=1, i<1000, Increment(i))
# Note using 10^3 is much slower. Mma3.0 also is slower with 10^3

function do_For(mx::Mxpr{:For}, start, test, incr)
    @unsetbreak
    doeval(start)
    while
        (res = doeval(test))
        @checkbreak res
        res = doeval(incr)
        @checkbreak res
    end
    return Null
end

function do_For(mx::Mxpr{:For}, start, test, incr, body)
    @unsetbreak
    doeval(start)
    while (res = doeval(test))
        @checkbreak res
        res = doeval(body)
        @checkbreak res
        @checkthrowreturn(res)
        @checkcontinue(res, incr)
        doeval(incr)
    end
    return Null
end

### If

@mkapprule If

@sjdoc If """
    If(test, tbranch, fbranch)

evaluate test and if the result is `True`, evaluate `tbranch`, otherwise `fbranch`.

    If(test, tbranch)

return `Null` if `test` does not evaluate to `True`.

    If(test, tbranch, fbranch, ubranch)

evaluate `ubranch` if the truth value of `test` cannot be determined.
"""

do_If(mxpr::Mxpr{:If}, test, tbranch) =  doeval(test) == true ? doeval(tbranch) : Null

function do_If(mxpr::Mxpr{:If}, test, tbranch, fbranch)
    tres = doeval(test)
    return tres == true ? doeval(tbranch) : tres == false ? doeval(fbranch) : mxpr
end

function do_If(mxpr::Mxpr{:If}, test, tbranch, fbranch, ubranch)
    tres = doeval(test)
    return tres == true ? doeval(tbranch) : tres == false ? doeval(fbranch) : doeval(ubranch)
end

### While

@sjdoc While """
    While(test, body)

evaluates `test` then `body` in a loop until `test` does not return `True`.
"""

@mkapprule While :nargs => 1:2

# TODO: Check for return and continue here too
function do_While(mx::Mxpr{:While}, test)
    @unsetbreak
    while (res = doeval(test)) == true
        @checkbreak res
    end
    return Null
end

function do_While(mx::Mxpr{:While}, test, body)
    @unsetbreak
    while doeval(test) == true
        res = doeval(body)
        @checkbreak res
        @checkthrowreturn(res)
        @checkcontinue0(res)
    end
    return Null
end

### Break

@sjdoc Break """
    Break()

exit the nearest enclosing `For`, `While`, or `Do` loop.
"""

@mkapprule Break

function do_Break(mx::Mxpr{:Break})
    set_break()
    return Null
end

#### Return

@sjdoc Return """
    Return(x)

return `x` from the nearest enclosing block.

    Return()

return `Null`.

Calling `Return(x)` from within `For`, `While`, `CompoundExpression` returns `Return(x)`.
Calling `Return(x)` from within `Do` and `Module` returns `x`.
"""

### Continue

@sjdoc Continue """
    Continue()

begin the next iteration of the enclosing loop without evaluating any remaining
expressions in the body.
"""

### Do

@sjdoc Do """
    Do(expr, [imax])

evaluate `expr` `imax` times.

    Do(expr, [i, imax])

evaluate `expr` `imax` times with `i` localized and taking values from `1` through
`imax` in increments of `1`.

    Do(expr, [i, imin, imax])

evaluate `expr` with `i` taking values from `imin` to `imax` with increment `1`.
`imin` and `imax` may be symbolic.

    Do(expr, [i, imin, imax, di])

evaluate `expr` with `i` taking values from `imin` to `imax` with increment `di`.
`imin`, `imax`, and `di` may be symbolic.

    Do(expr, [i, [i1, i2, ...])

evaluate `expr` with `i` taking values `i1, i2, ...`.

The binding of variable `i` has lexical scope, as in `Module`.
"""

function apprules(mx::Mxpr{:Do})
    expr = mx[1]
    iter = make_sjiter(mx[2])
    do_doloop(expr, iter)
end

function do_doloop(expr, iter::SJIter1)
    do_doloop_kern(expr, iter.imax)
end

# TODO: prbly don't need to use kernel
function do_doloop_kern(expr, imax)
    start = one(imax)
    @unsetbreak
    for i in start:imax
        res = doeval(expr)
        @checkbreak res
        @checkreturn(res)
        @checkcontinue0(res)
    end
    Null
end

function do_doloop(expr, iter::SJIter2)
    isym = get_localized_symbol(iter.i)
    ex = substitute_symbol(deepcopy(expr), iter.i, isym)
    @unsetbreak
    for i in 1:iter.imax  # mma makes i an Int no matter the type of iter.imax
        setsymata(isym, i)
        res = doeval(ex)
        @checkbreak res
        @checkreturn(res)
        @checkcontinue0(res)
    end
    delete_sym(isym)
    Null
end

function do_doloop(expr, iter::SJIter3{T,V}) where {T<:Real,V<:Real}
    isym = get_localized_symbol(iter.i)
    ex = substitute_symbol(deepcopy(expr), iter.i, isym)
    @unsetbreak
    for i in iter.imin:iter.imax  # mma makes i type of one of these
        setsymval(isym, i)
        res = doeval(ex)
        @checkbreak res
        @checkreturn(res)
        @checkcontinue0(res)
    end
    delete_sym(isym)
    Null
end

# fields of iter may be symbolic
function do_doloop(expr, iter::SJIter3)
    isym = get_localized_symbol(iter.i)
    ex = substitute_symbol(deepcopy(expr), iter.i, isym)
    setsymval(isym, iter.imin)
    @unsetbreak
    for i in 1:(iter.num_iters)
        res = doeval(ex)
        @checkbreak res
        @checkreturn(res)
        @checkcontinue0(res)  # This will not increment isym
        setsymval(isym, doeval(mxpr(:Plus, isym, 1)))
    end
    delete_sym(isym)
    Null
end

function do_doloop(expr, iter::SJIter4{T,V,W}) where {T<:Real, V<:Real, W<:Real}
    isym = get_localized_symbol(iter.i)
    ex = substitute_symbol(deepcopy(expr), iter.i, isym)
    @unsetbreak
    for i in (iter.imin):(iter.di):(iter.imax)
        setsymval(isym, i)
        res = doeval(ex)
        @checkbreak res
        @checkreturn(res)
        @checkcontinue0(res)
    end
    delete_sym(isym)
    Null
end

# fields of iter may be symbolic
function do_doloop(expr, iter::SJIter4)
    isym = get_localized_symbol(iter.i)
    ex = substitute_symbol(deepcopy(expr), iter.i, isym)
    setsymval(isym, iter.imin)
    @unsetbreak
    for i in 1:(iter.num_iters)
        res = doeval(ex)
        @checkbreak res
        @checkreturn(res)
        @checkcontinue0(res)
        setsymval(isym, doeval(mplus(isym, iter.di)))
    end
    delete_sym(isym)
    Null
end

function do_doloop(expr, iter::SJIterList)
    isym = get_localized_symbol(iter.i)
    ex = substitute_symbol(deepcopy(expr), iter.i, isym)
    @unsetbreak
    for i in 1:(length(iter.list))
        setsymval(isym, iter.list[i])
        res = doeval(ex)
        @checkbreak res
        @checkreturn(res)
        @checkcontinue0(res)
    end
    delete_sym(isym)
    Null
end

### CompoundExpression

@sjdoc CompoundExpression """
    CompoundExpression(expr1, expr2, ...) or (expr1, expr2, ...)

evaluate each expression in turn and return the result of only the final evaluation.
"""
function apprules(mx::Mxpr{:CompoundExpression})
    local res
        @inbounds for i in 1:length(mx)
            res = doeval(mx[i])
            is_throw() && return res
            is_break() && break
            @checkthrowreturn(res)
            @checkthrowcontinue(res)
        end
    res
end

### Warn

@mkapprule Warn :nodefault => true

@doap Warn(args...) = @symwarn(args...)

### Throw

@mkapprule Throw :nargs => 1:2

@doap function Throw(x)
    set_throw()
    return mx
end

@doap function Throw(x, u)
    set_throw()
    return mx
end

### Catch

# TODO: implement tagged throw, catch
@mkapprule Catch :nargs => 1:2

@doap function Catch(x)
    res = doeval(x)
    if is_throw()
        if is_Mxpr(res, :Throw) && length(res) == 1
            clear_throw()
            return length(res) == 0 ? Null : res[1]
        end
        @symwarn("Catch: throw set but expr is ", res)
        return res
    end
    return res
end

@doap function Catch(x, tag)
    res = doeval(x)
    if is_throw()
        if is_Mxpr(res, :Throw) && length(res) == 2 && matchq(res[2], tag)
            clear_throw()
            return length(res) == 0 ? Null : res[1]
        else
            return res
        end
    end
    return res
end


text:
struct ProblemResultsExport
    name::String
    duals::Set{Symbol}
    parameters::Set{Symbol}
    variables::Set{Symbol}
    optimizer_stats::Bool

    function ProblemResultsExport(name, duals, parameters, variables, optimizer_stats)
        duals = _check_fields(duals)
        parameters = _check_fields(parameters)
        variables = _check_fields(variables)
        new(name, duals, parameters, variables, optimizer_stats)
    end
end

function ProblemResultsExport(
    name::AbstractString;
    duals = Set{Symbol}(),
    parameters = Set{Symbol}(),
    variables = Set{Symbol}(),
    optimizer_stats = true,
)
    return ProblemResultsExport(name, duals, parameters, variables, optimizer_stats)
end

function _check_fields(fields)
    if !(typeof(fields) <: Set{Symbol})
        fields = Set(Symbol.(fields))
    end

    if :all in fields && length(fields) > 1
        throw(IS.InvalidValue("'all' can only be present if the array has one element"))
    end

    return fields
end

should_export_dual(x::ProblemResultsExport, name) = _should_export(x, :duals, name)
should_export_parameter(x::ProblemResultsExport, name) =
    _should_export(x, :parameters, name)
should_export_variable(x::ProblemResultsExport, name) = _should_export(x, :variables, name)

function _should_export(exports::ProblemResultsExport, field_name, name)
    container = getfield(exports, field_name)
    isempty(container) && return false
    first(container) == :all && return true
    return name in container
end