text:
module CAS

export SubDeg, SubZero, SubAll, SubEven
export sym_multivector

using SymbolicUtils
using ArgCheck
using Accessors
using OffsetArrays
using SymbolicUtils.Code

################################################################################
##### BasisBlade
################################################################################
struct BasisBlade
    # Given a basis e1,e2,e3,e4,e5
    #
    # The 5*e214 would be encoded as follows:
    # First normalize to sorted wedge: 5*e214 =  -5 *e124
    # coeff = -5
    # mask = [true, true, false, true, false] # no e3,e5
    coeff
    mask
end

function check_compatible(b1::BasisBlade, b2::BasisBlade)
    @argcheck axes(b1.mask) == axes(b2.mask)
end

function Base.isequal(b1::BasisBlade, b2::BasisBlade)
    isequal(b1.coeff, b2.coeff) && isequal(b1.mask, b2.mask)
end

function Base.:(==)(b1::BasisBlade, b2::BasisBlade)
    ==(b1.coeff, b2.coeff) && ==(b1.mask, b2.mask)
end

function genindices(o::BasisBlade)
    axes(o.mask,1)
end

function degree(o::BasisBlade)
    sum(o.mask)
end

function zerobasisblade(r::AbstractUnitRange)
    BasisBlade(0, fill(false, r))
end

function Base.:(-)(b::BasisBlade)
    @set b.coeff *= -1
end

function count_transpositions(mask1, mask2)
    @argcheck axes(mask1) == axes(mask2)
    transpositions = 0
    disjoint = true
    n = lastindex(mask1)
    for i in eachindex(mask1, mask2)
        if mask2[i]
            if mask1[i]
                disjoint = false
            end
            for j in (i+1):n
                if mask1[j]
                    transpositions += 1
                end
            end
        end
    end
    return (;transpositions, disjoint)
end

function wedge(b1::BasisBlade, b2::BasisBlade)
    @argcheck genindices(b1) == genindices(b2)
    intel = count_transpositions(b1.mask, b2.mask)
    if intel.disjoint
        coeff = (-1)^(intel.transpositions) * b1.coeff * b2.coeff
        mask::Mask = map(xor, b1.mask, b2.mask)
        ret = BasisBlade(coeff, mask)
        @check degree(ret) == degree(b1) + degree(b2)
        return ret
    else
        return zerobasisblade(genindices(b1))
    end
end

function geomul(b1::BasisBlade, b2::BasisBlade, metric::AbstractVector)
    @argcheck axes(b1.mask) == axes(b2.mask) == axes(metric)

    mask = fill(false, genindices(b1))
    coeff = b1.coeff * b2.coeff
    mask1 = b1.mask
    mask2 = b2.mask
    for i in eachindex(mask)
        if (mask1[i]) && (mask2[i])
            coeff *= metric[i]
            mask[i] = false
        elseif (mask1[i] && (!mask2[i]))
            mask[i] = true
        elseif (!mask1[i]) && mask2[i]
            mask[i] = true
        elseif (!mask1[i]) && (!mask2[i])
            mask[i] = false
        else
            error("Unreachable")
        end
    end
    sign = (-1)^count_transpositions(mask1, mask2).transpositions
    coeff *= sign
    return BasisBlade(coeff, mask)
end

function inner(b1::BasisBlade, b2::BasisBlade, metric)
    ret = geomul(b1, b2, metric)
    g1 = degree(b1)
    g2 = degree(b2)
    if degree(ret) == abs(g1 - g2)
        return ret
    else
        return zerobasisblade(eachindex(b1.mask))
    end
end

function anti_wedge(b1::BasisBlade, b2::BasisBlade)
    dual(wedge(dual(b1), dual(b2)))
end

function rev(b::BasisBlade)
    sign = (-1)^(degree(b)÷2)
    return BasisBlade(sign * b.coeff, b.mask)
end

function dual(b1::BasisBlade)
    # TODO figure out good sign convention for dual
    @assert genindices(b1) == 0:3
    d = degree(b1)
    mask::Mask = map((!), b1.mask)
    ntrans = if sum(b1.mask) <= sum(mask)
        count_transpositions(b1.mask, mask).transpositions
    else
        count_transpositions(mask, b1.mask).transpositions
    end
    sign = (-1)^ntrans
    return BasisBlade(sign*b1.coeff, mask)
end

function create_mask(r::AbstractUnitRange, inds...)
    @argcheck issorted(inds)
    mask = fill(false,r)
    for i in inds
        mask[i] = true
    end
    return mask
end

function add(b1::BasisBlade, b2::BasisBlade)
    @argcheck b1.mask == b2.mask
    BasisBlade(b1.coeff + b2.coeff, b1.mask)
end

################################################################################
##### MultiVector
################################################################################

const Mask = OffsetVector{Bool, Vector{Bool}}

struct MultiVector
    summands::Dict{Mask, Any}
    metric::OffsetVector{Float64, Vector{Float64}}
end

function check_compatible(m1::MultiVector, m2::MultiVector)
    @argcheck m1.metric == m2.metric
end

const ∧ = wedge
const ∨ = anti_wedge

function optimize!(m::MultiVector)
    for (key,val) in pairs(m.summands)
        m.summands[key] = simplify(val, expand=true)
    end
    prunezeros!(m)
end

function prunezeros!(m::MultiVector)
    filter!(m.summands) do (key, val)
        if val isa Number
            !iszero(val)
        else
            true
        end
    end
    m
end

function operate_on_basis(f, m::MultiVector)
    res = Dict{Mask, Any}()
    for (mask, coeff) in pairs(m.summands)
        b = f(BasisBlade(coeff, mask))
        if haskey(res, b.mask)
            res[b.mask] += b.coeff
        else
            res[b.mask] = b.coeff
        end
    end
    return optimize!(MultiVector(res, m.metric))
end

function operate_on_basis(f, m1::MultiVector, m2::MultiVector)
    check_compatible(m1, m2)
    res = Dict{Mask, Any}()
    for (mask1, coeff1) in pairs(m1.summands)
        for (mask2, coeff2) in pairs(m2.summands)
            b = f(BasisBlade(coeff1, mask1), BasisBlade(coeff2, mask2))
            if haskey(res, b.mask)
                res[b.mask] += b.coeff
            else
                res[b.mask] = b.coeff
            end
        end
    end
    return optimize!(MultiVector(res, m1.metric))
end

function MultiVector(b::BasisBlade, metric)
    d = Dict(b.mask => b.coeff)
    optimize!(MultiVector(d, metric))
end
function mergewith_multivectors(f, ms...)
    m1 = first(ms)
    check_compatible(ms...)
    d = mergewith(f, map(m -> m.summands, ms)...)
    optimize!(MultiVector(d, m1.metric))
end

for op in [
        :(Base.:(+)),
        :(Base.:(-)),
    ]
    @eval function $(op)(m1::MultiVector, m2::MultiVector)
        mergewith_multivectors($op, m1, m2)
    end
end

for f in [:dual, :rev, :(Base.:(-))]
    @eval $f(m::MultiVector) = operate_on_basis($f, m)
end

for op in [:wedge, :anti_wedge]
    @eval $op(m1::MultiVector, m2::MultiVector) = operate_on_basis($op, m1, m2)
end
function inner(m1::MultiVector, m2::MultiVector)
    check_compatible(m1, m2)
    operate_on_basis(m1, m2) do b1, b2
        inner(b1, b2, m1.metric)
    end
end

function geomul(m1::MultiVector, m2::MultiVector)
    check_compatible(m1, m2)
    operate_on_basis(m1, m2) do b1, b2
        geomul(b1, b2, m1.metric)
    end
end

function project(onto::MultiVector, obj::MultiVector)
    geomul(inner(obj, onto), onto)
end

function Base.:(*)(m1::MultiVector, m2::MultiVector)
    geomul(m1, m2)
end
function scalar_mul(s, m2::MultiVector)
    d = Dict(key => s*val for (key, val) in pairs(m2.summands))
    MultiVector(d, m2.metric)
end
function Base.:(*)(s::Union{SymbolicUtils.Symbolic{Real}, Real}, m::MultiVector)
    scalar_mul(s, m)
end
function Base.:(*)(m::MultiVector, s::Real)
    scalar_mul(s, m)
end

#operate_on_basis(dual, m)

function Vector0(itr)
    OffsetArray(itr, (0:length(itr)-1))
end
const METRIC_PGA3D = Vector0([0.,1,1,1])

function sandwich(x::MultiVector,g::MultiVector)
    rev(g)*x*g
end

function Base.show(io::IO, x::MultiVector)
    pieces = String[]
    d = sort(x.summands)
    for (mask, val) in pairs(d)
        inds = eachindex(mask)[mask]
        push!(pieces, "($val)*e$(inds...)")
    end
    println(io, join(pieces, " + "))
end

################################################################################
##### Code gen
################################################################################

function basisblade_symbol_from_mask(prefix, mask)
    inds = eachindex(mask)[mask]
    Symbol(prefix, inds...)
end
basisblade_symbol_from_mask(mask) = basisblade_symbol_from_mask(:e, mask)

function sym_getproperty_from_mask(mask)
    ei = basisblade_symbol_from_mask(mask)
    S = SymbolicUtils.Sym{SymbolicUtils.FnType{Tuple{Any},Real}}
    return S(ei)
end

struct SubAll end
struct SubEven end
struct SubDeg{N} end
SubDeg(N) = SubDeg{N}()
struct SubZero end

const SUBS = [SubAll(),
              SubEven(),
              SubDeg{0}(),
              SubDeg{1}(),
              SubDeg{2}(),
              SubDeg{3}(),
              SubDeg{4}(),
              SubZero(),
             ]

function masks_inside(::SubEven)
    ft = (false, true)
    [Vector0([a,b,c,d]) for a in ft for b in ft for c in ft for d in ft if iseven(a+b+c+d)]
end

function masks_inside(::SubDeg{N}) where {N}
    ft = (false, true)
    [Vector0([a,b,c,d]) for a in ft for b in ft for c in ft for d in ft if a+b+c+d == N]
end

function masks_inside(::SubAll)
    ft = (false, true)
    [Vector0([a,b,c,d]) for a in ft for b in ft for c in ft for d in ft]
end
function masks_inside(::SubZero)
    []
end

function sym_multivector(symbol_type::Pair)
    symbol, sub = symbol_type
    d = Dict(
        mask => begin
            xijk::Symbol = basisblade_symbol_from_mask(symbol, mask)
            SymbolicUtils.Sym{Real}(xijk)
        end
        for mask in masks_inside(sub)
    )
    metric = METRIC_PGA3D
    return MultiVector(d, metric)
end

function expr_from_arg(arg::Pair)
    argname, TODO = arg
    argname
end

function symbol_MultiVectorSub(sub)
    @argcheck !(sub isa DataType)
    if sub isa SubZero
        :MultiVectorZero
    elseif sub isa SubDeg{0}
        :MultiVector0
    elseif sub isa SubDeg{1}
        :MultiVector1
    elseif sub isa SubDeg{2}
        :MultiVector2
    elseif sub isa SubDeg{3}
        :MultiVector3
    elseif sub isa SubDeg{4}
        :MultiVector4
    elseif sub isa SubEven
        :MultiVectorEven
    elseif sub isa SubAll
        :MultiVectorAll
    else
        error("Unreachable sub = $sub")
    end
end

function expr_structdef(sub)
    if sub isa SubZero
        return :(struct MultiVectorZero end)
    end
    MV = symbol_MultiVectorSub(sub)
    masks = masks_inside(sub)
    T = :T
    eijs = map(basisblade_symbol_from_mask, masks)
    fields = map(eijs) do eij
        :($eij::$T)
    end
    _eijs= map(masks) do mask
        basisblade_symbol_from_mask("_e", mask)
    end
    eijsT = map(eijs) do eij
        :(convert($T, $eij))
    end
    ctor1 = :(
              function $MV{$T}($(eijs...)) where {$T}
                  new{$T}($(eijsT...))
              end
    )
    ctor2 = :(
        function $MV($(eijs...))
            $(Expr(:tuple, _eijs...)) = promote($(eijs...))
            T = typeof($(first(_eijs)))
            new{$T}($(_eijs...))
        end
    )
    ret = Expr(:struct, false, :($MV{$T}), Expr(:block, fields..., ctor1, ctor2))
    Base.remove_linenums!(ret)
    return ret
end

function exprs_structdef_batteries(sub)
    MV = symbol_MultiVectorSub(sub)
    masks = masks_inside(sub)
    ret = []
    push!(ret, expr_structdef(sub))
    #if sub isa SubZero
    #    return ret
    #end
    kws = map(masks) do mask
        eij = basisblade_symbol_from_mask(mask)
        Expr(:kw, eij, 0)
    end
    if sub != SubZero()
        kwconstructor = :(
          function $MV(;$(kws...))
              return $MV($(map(basisblade_symbol_from_mask, masks)...))

          end
         )
        push!(ret, kwconstructor)
    end

    N = length(masks)
    tuple_access = [:(t[$i]) for i in 1:N]
    from_Tuple = :(
        function $MV(t::NTuple{$N,Any})
            $MV($(tuple_access...))
        end
    )
    push!(ret, from_Tuple)

    field_accesses = map(masks) do mask
        eij = basisblade_symbol_from_mask(mask)
        :(x.$eij)
    end
    to_Tuple = :(
        function Base.Tuple(x::$MV)
            tuple($(field_accesses...))
        end
    )
    push!(ret, to_Tuple)

    nb = length(masks_inside(sub))
    def_basislength = :(
        function basislength(::Type{<:$MV})
            $nb
        end
    )
    push!(ret, def_basislength)
    if sub !== SubZero()
        def_scalartype = :(
        function scalartype(::Type{$MV{T}}) where {T}
            T
        end
        )
        push!(ret, def_scalartype)
        def_ctor = :(
            function $MV{T}(x::$MV) where {T}
                $MV{T}($(field_accesses...))
            end
        )
        push!(ret, def_ctor)
    end

    foreach(Base.remove_linenums!, ret)
    ret
end

function expr_fundef(f_symbolic, f_name, args::Pair...)
    notation = []
    arg_multivectors = []
    arg_exprs = []
    for (argsymbol, sub) in args
        d = Dict()
        for mask in masks_inside(sub)
            argsymbol_ = Symbol(argsymbol, "_")
            lhs = basisblade_symbol_from_mask(argsymbol_, mask)
            d[mask] = SymbolicUtils.Sym{Real}(lhs)
            rhs = :($(argsymbol).$(basisblade_symbol_from_mask(:e, mask)))
            push!(notation, :($lhs = $rhs))
        end
        mv = MultiVector(d, METRIC_PGA3D)
        push!(arg_multivectors,mv)
        T = symbol_MultiVectorSub(sub)
        push!(arg_exprs, :($argsymbol::$T))
    end
    body = []
    append!(body, notation)
    sym_ret::MultiVector = f_symbolic(arg_multivectors...)
    notation_ret = []
    ret_masks = sort!(collect(keys(sym_ret.summands)))
    for mask in ret_masks
        sym = sym_ret.summands[mask]
        lhs = basisblade_symbol_from_mask(:ret_, mask)
        rhs = toexpr(sym)
        ex = :($lhs = $rhs)
        push!(notation_ret, ex)
    end
    append!(body, notation_ret)
    sub = minimal_sub_containing_masks(ret_masks)
    ctor = symbol_MultiVectorSub(sub)

    kws = map(ret_masks) do mask
        # eij = retij
        Expr(:kw,
                basisblade_symbol_from_mask(:e, mask),
                basisblade_symbol_from_mask(:ret_, mask)
        )
    end

    construct_ret = Expr(:call, ctor, Expr(:parameters, kws...))
    if sub isa SubZero
        empty!(body)
    end
    push!(body, :(return $construct_ret))
    ex = :(function $f_name($(arg_exprs...)) $(body...) end)
    Base.remove_linenums!(ex)
    ex
end

function minimal_sub_containing_masks(masks::AbstractArray)
    grades = unique(map(sum, masks))
    if length(grades) == 0
        SubZero()
    elseif length(grades) == 1
        SubDeg{only(grades)}()
    elseif all(iseven, grades)
        SubEven()
    else
        SubAll()
    end
end

function exprs_everything()
    ret = []
    for sub in SUBS
        exs = exprs_structdef_batteries(sub)
        append!(ret, exs)
    end
    for (f, fsym) in [
            (dual, :dual),
            (rev, :rev),
            (-, :(Base.:-)),

        ]
        for sub in SUBS
            ex = expr_fundef(f, fsym, :x => sub)
            push!(ret, ex)
        end
    end
    for (op, opsym) in [
            (+, :(Base.:+)),
            (-, :(Base.:-)),
            (geomul, :geomul),
            (wedge, :wedge),
            (sandwich, :sandwich),
            (anti_wedge, :anti_wedge),
            (inner, :inner),

        ]
        for sub1 in SUBS
            for sub2 in SUBS
                ex = expr_fundef(op, opsym, :x => sub1, :y => sub2)
                push!(ret, ex)
            end
        end
    end
    GRADE_SUBS = [SubDeg(0), SubDeg(1), SubDeg(2), SubDeg(3), SubDeg(4)]
    for sub1 in GRADE_SUBS
        for sub2 in GRADE_SUBS
            ex = expr_fundef(project, :project, :onto=>sub1, :obj => sub2)
            push!(ret, ex)
        end
    end
    # conversion
    for sub in SUBS
        sub == SubZero() && continue
        ex = expr_convert(SubAll(), sub)
        push!(ret, ex)
    end
    for sub in SUBS
        ex = expr_convert(sub, SubZero())
        push!(ret, ex)
    end
    for sub in [SubDeg(0), SubDeg(2), SubDeg(4)]
        ex = expr_convert(SubEven(), sub)
        push!(ret, ex)
    end
    return ret
end

function expr_convert(sub_big, sub_small)
    MV_big = symbol_MultiVectorSub(sub_big)
    MV_small = symbol_MultiVectorSub(sub_small)
    masks = masks_inside(sub_small)
    kws = map(masks) do mask
        eij = basisblade_symbol_from_mask(mask)
        Expr(:kw, eij, :(x.$eij))
    end
    ex = :(
        function Base.convert(MV::Type{<:$(MV_big)}, x::$MV_small)
            MV($MV_big(;$(kws...)))
        end
    )
    Base.remove_linenums!(ex)
    ex
end

function print_everything(io::IO)
    for ex in exprs_everything()
        s = string(ex)
        println(io, s)
    end
end

end#module


text:
#=
jim.jl
jiffy image display
2019-02-23 Jeff Fessler, University of Michigan
=#

export jim

import MIRTjim # jim

@deprecate jim(args... ; kwargs...) MIRTjim.jim(args... ; kwargs...)

ext: jl


text:
module mtmodel

using Flux, ArgCheck, StatsBase, YAML
using Flux: gate
using Flux.Tracker: istracked

using ..modelutils
import ..modelutils: load!, randn_repar, mlp, get_strtype_wo_params
export load!

const D_IM = 32

################################################################################
##                                                                            ##
##              Define a basic MTGRU model (with no inputs)                   ##
##                                                                            ##
################################################################################

"""
    batch_matvec(A::Tensor, X::Matrix)

for A ∈ ℝ^{n × m × d}, X ∈ ℝ^{n × d}. Performs ``n`` matrix-vec multiplications

    A[i,:,:] * X[i,:]

for ``i \\in 1,\\ldots,n``. This is performed via expansion of `X` and can be
efficiently evaluated using BLAS operations, which (I have assumed!) will often
be faster than a loop, especially in relation to AD and GPUs.
"""
batch_matvec(A::AbstractArray{T,3}, X::AbstractArray{T,2}) where T =
    dropdims(sum(A .* unsqueeze(X, 1), dims=2), dims=2)

Flux.gate(x::AbstractArray{T,3}, h, n) where T = x[gate(h,n),:,:]
Flux.gate(x::AbstractArray{T,4}, h, n) where T = x[gate(h,n),:,:,:]

mutable struct MTGRU_NoU{F}
    N::Int
    G::F
end

"""
    MTGRU_NoU(N, G)
Multi-task GRU (Gated Recurrent Unit) model. Produces a GRU layer
for each value of `z` ∈ Rᵈ which depends on the map `G` inside
the MTGRU. N = out dimension, G: Rᵈ->R^{nLSTM}. z can be batched
as `z` ∈ R^{d × nbatch}.

In order to simplify the implementation, there is no input depen-
dent evolution (hence no `U`s, and hence no `Wi`). It is
straight-forward to extend to this case, should it be reqd.
"""
function (m::MTGRU_NoU)(z, h=nothing)
    N, nB = m.N, size(z, 2)
    λ = m.G(z)
    Wh = reshape(λ[1:N*N*3,:], N*3, N, nB)
    b = λ[N*N*3+1:(N+1)*N*3, :]
    open_forget_gate = zero(b)   # no backprop, and hence not tracked, even if `b` is.
    open_forget_gate[gate(N, 2), :] .= 1
    b += open_forget_gate
    h = something(h, istracked(λ) ? Flux.param(zero(b)) : zero(b))

    Flux.Recur(BatchedGRUCell_NoU(Wh, b, h))
end

Flux.@treelike MTGRU_NoU
Base.show(io::IO, l::MTGRU_NoU) =
  print(io, "Multitask-GRU(", l.N, ", ", typeof(l.G), ", no inputs)")

"""
  MTGRU_NoU_fixb(N, G, b)
Multi-task GRU (Gated Recurrent Unit) model. Produces a GRU layer
for each value of `z` ∈ Rᵈ which depends on the map `G` inside
the MTGRU. N = out dimension, G: Rᵈ->R^{nLSTM}. z can be batched
as `z` ∈ R^{d × nbatch}. This version of the `MTGRU_NoU` does not
adapt the offset vector `b` in the GRU, but learns a fixed value
for all `z`.

In order to simplify the implementation, there is no input depen-
dent evolution (hence no `U`s, and hence no `Wi`). It is
straight-forward to extend to this case, should it be reqd.
"""
mutable struct MTGRU_NoU_fixb{F,V}
  N::Int
  G::F
  b::V
end

"""
  MTGRU_NoU_fixb_v2(N, G, b)
As per `MTGRU_NoU_fixb`, except now only the *state* bias is fixed across all
tasks. The gate biases depend on `z` just like the `MTGRU_NoU`.
"""
mutable struct MTGRU_NoU_fixb_v2{F,V}
  N::Int
  G::F
  b::V
end

function (m::MTGRU_NoU_fixb)(z, h=nothing)
  N, nB = m.N, size(z, 2)
  λ = m.G(z)
  Wh = reshape(λ, N*3, N, nB)
  b = m.b
  open_forget_gate = zero(b)   # no backprop, and hence not tracked, even if `b` is.
  open_forget_gate[gate(N, 2), :] .= 1
  b += open_forget_gate
  h = something(h, istracked(λ) ? Flux.param(zero(b)) : zero(b))

  Flux.Recur(BatchedGRUCell_NoU(Wh, b, h))
end

function (m::MTGRU_NoU_fixb_v2)(z, h=nothing)
    N, nB = m.N, size(z, 2)
    λ = m.G(z)
    Wh = reshape(λ[1:N*N*3,:], N*3, N, nB)
    gmove = Flux.has_cuarrays() && Tracker.data(z) isa Flux.CuArray ? gpu : identity
    expander = gmove(ones(eltype(Tracker.data(λ)), 1, nB))
    b = vcat(λ[N*N*3+1:N*N*3+N*2, :], m.b * expander)
    open_forget_gate = zero(b)
    open_forget_gate[gate(N, 2), :] .= 1
    b += open_forget_gate
    h = something(h, istracked(λ) ? Flux.param(zero(b)) : zero(b))

    Flux.Recur(BatchedGRUCell_NoU(Wh, b, h))
end

Flux.@treelike MTGRU_NoU_fixb
Base.show(io::IO, l::MTGRU_NoU_fixb) =
print(io, "Multitask-GRU-fixb(", l.N, ", ", typeof(l.G), ", bias vector: ", typeof(l.b), "; no inputs)")

Flux.@treelike MTGRU_NoU_fixb_v2
Base.show(io::IO, l::MTGRU_NoU_fixb_v2) =
print(io, "Multitask-GRU-fixb_v2(", l.N, ", ", typeof(l.G), ", bias vector (state only): ", typeof(l.b), "; no inputs)")

"""
    BatchedGRUCell_NoU(Wh, b, h)
Multi-task GRU Cell (which takes no input).
`Wh`, `b`, `h` are the (concatenated) transformation matrices,
offsets, and initial hidden state respectively.

Each time the cell is called (no arguments) it will perform a one-
step evolution of the hidden state and return the current value.
The cell is implemented in batch-mode, and the final dimension of
each quantity is the batch index.
"""
mutable struct BatchedGRUCell_NoU{A,W}
    Wh::A
    b::W
    h::W
end

function (m::BatchedGRUCell_NoU)(h, x=nothing)
  b, o = m.b, size(h, 1)
  gh = batch_matvec(m.Wh, h)
  r = σ.(gate(gh, o, 1) .+ gate(b, o, 1))
  z = σ.(gate(gh, o, 2) .+ gate(b, o, 2))
  h̃ = tanh.(r .* gate(gh, o, 3) .+ gate(b, o, 3))
  h′ = (1 .- z).*h̃ .+ z.*h
  return h′, h′
end

Flux.@treelike BatchedGRUCell_NoU
Base.show(io::IO, l::BatchedGRUCell_NoU) =
  print(io, "Batched GRU Cell(dim=", size(l.Wh,2), ", batch=", size(l.Wh,3), ", no inputs)")
Flux.hidden(m::BatchedGRUCell_NoU) = m.h

################################################################################
##                                                                            ##
##           Define *THE* multi-task model used in our experiments            ##
##                                                                            ##
################################################################################

struct MTSeq_CNN end
struct MTSeq_Single end
struct MTSeq_Double end

"""
    MTSeqModel_Pred(enc_tform, init_enc, init_post, mt_enc, mt_post,
        x0decoder, mtbias_deocder, gen_model, d_y, mtbias_only)

A sequence-to-sequence model 𝒴→𝒴 which imbibes `T_enc` steps of the sequence,
and predicts `T_steps` (including the original `T_enc`), but unlike the
`Seq2SeqModel`, the `MTSeqModel_Pred` can *modulate* its prediction based on
hierarchical 'multi-task' variables. The `Pred` suffix refers to the function
of the encoders. I had previously made the `E3` suffix version (see docstring),
but this was optimized for learning a latent representation of the entire
sequence, which is effective for disentangled-like representation and likelihood
but less so for prediction, since optimizing the evidence of the **observed**
sequence appears to perform worse on future observations than an RNN trained
only to predict. For this model, the encoder learns a latent code that is
optimized for **predictive** performance, not **reconstructive** performance.
See the paper, and similar observations have previously been made -- see e.g.
Karl et al. 2016, Deep VB Filter.

As for the `Seq2SeqModel`, the posterior over the initial state of the generator
is learned variationally via the `init_enc` RNN with the final state passed to
the `init_post` layer which generates the mean and variance of the latent
posterior. Since the posterior is typically lower dimensional than the state of
the generator network, the latent variable is expanded via the `x0decoder`
network into the state of the `generator`. The MT variable is encoded with a
full-length sequence encoding via the `mt_enc` network, and the final state is
passed to the `mt_post` layer to generate the posterior mean and variance.
The MT variable enters the generator network either:

1. by predicting all parameters of the network via a `MTGRU_NoU` construction
(see docstring).
2. by forming a latent (constant) input at each time step, which functions as a
variable bias of the generator recurrent cell. This may be directly from the
latent variables, or via a decoder (`mtbias_decoder`).

Once these variables are sampled, the `generator` is iterated in an open-loop
fashion for `T_steps`, similar to the `Seq2SeqModel`. The penultimate slot `d`
is a record of the dimension of 𝒴, used currently just for `show`; and the
`mtbias_only` variable dictates whether either approach (1) or (2) is used as
above.

--------------------
This type can be called via
```julia
    m = MTSeqModel_Pred(...)
    m(y::Vector)
```
where `y` is a vector of slices of the sequence for the encoder. This means each
`y[i]` is an `AbstractMatrix` of size ``d \\times n_{\\text{batch}}`.

Options include `T_enc`, `T_steps` (already discussed), `reset` (which performs
a state reset before execution; true by default), and `stoch`, whether to
randomly sample the latent state, or use the state mean (default true).
"""
struct MTSeqModel_Pred{U,A,B,V,W,S,N,M}
    enc_tform::U
    init_enc::A
    init_post::B
    mt_enc::A
    mt_post::B
    x0decoder::V
    mtbias_decoder::W
    generator::S
    dec_tform::N
    model_type::M
    d::Int64
    mtbias_only::Bool
end

"""
    MTSeqModel_E3(enc_tform, init_enc, init_post, mt_enc, mt_post, chaos_enc, chaos_post,
        x0decoder, mtbias_deocder, gen_model, d_y, mtbias_only)

A sequence-to-sequence model 𝒴→𝒴 which imbibes `T_enc` steps of the sequence,
and predicts `T_steps` (including the original `T_enc`), but unlike the
`Seq2SeqModel`, the `MTSeqModel_E3` can *modulate* its prediction based on
hierarchical 'multi-task' variables. (The `E3` suffix refers to the number of
encoders: previously I had `E1`, `E2` and `E3` types, but this seems redundant.)

As for the `Seq2SeqModel`, the posterior over the initial state of the generator
is learned variationally via the `init_enc` RNN with the final state passed to
the `init_post` layer which generates the mean and variance of the latent
posterior. Since the posterior is typically lower dimensional than the state of
the generator network, the latent variable is expanded via the `x0decoder`
network into the state of the `generator`. The MT variable is encoded with a
full-length sequence encoding via the `mt_enc` network, and the final state is
passed to the `mt_post` layer to generate the posterior mean and variance. A
similar set of operations is performed for the more local 'chaos variable'
using the `chaos_enc` and `chaos_post` networks. The 'chaos variable' encodes
sequence variation due to inherent unpredictability of a system, perhaps due to
chaos. This allows different long term predictions from the same state. The MT
and chaos variables enter the generator network either:

1. by predicting all parameters of the network via a `MTGRU_NoU` construction
(see docstring).
2. by forming a latent (constant) input at each time step, which functions as a
variable bias of the generator recurrent cell. This may be directly from the
latent variables, or via a decoder (`mtbias_decoder`).

Once these variables are sampled, the `generator` is iterated in an open-loop
fashion for `T_steps`, similar to the `Seq2SeqModel`. The penultimate slot `d`
is a record of the dimension of 𝒴, used currently just for `show`; and the
`mtbias_only` variable dictates whether either approach (1) or (2) is used as
above.

--------------------
This type can be called via
```julia
    m = MTSeqModel_E3(...)
    m(y::Vector, yfull::Vector)
```
where `y` is a vector of slices of the sequence for the encoder, and `yfull` is
the entire sequence (inc the future) on which to infer the latent MT variable
`z`. For both `y` and `yfull`, `y[i]` is an `AbstractMatrix` of size
``d \\times n_{\\text{batch}}`.

Options include `T_enc`, `T_steps` (already discussed), `reset` (which performs
a state reset before execution; true by default), and `stoch`, whether to
randomly sample the latent state, or use the state mean (default true).
"""
struct MTSeqModel_E3{U,A,B,V,W,S,N,M}
    enc_tform::U
    init_enc::A
    init_post::B
    mt_enc::A
    mt_post::B
    chaos_enc::A
    chaos_post::B
    x0decoder::V
    mtbias_decoder::W
    generator::S
    dec_tform::N
    model_type::M
    d::Int64
    mtbias_only::Bool
end

# Define union type for convenience
MTSeqModels{U,A,B,V,W,S,N,M} = Union{MTSeqModel_Pred{U,A,B,V,W,S,N,M}, MTSeqModel_E3{U,A,B,V,W,S,N,M}}

# Model utils
Flux.@treelike MTSeqModel_Pred
Flux.@treelike MTSeqModel_E3
model_type(m::MTSeqModels{U,A,B,V,W,S,N,M}) where {U,A,B,V,W,S,N,M} = M

unpack(m::MTSeqModel_Pred) = Flux.children(m)[1:9] # unpack struct; exclude `d`, `mtbias_only`.
unpack(m::MTSeqModel_E3) = Flux.children(m)[1:11]
unpack_inference(m::MTSeqModel_Pred) = Flux.children(m)[1:5]
unpack_inference(m::MTSeqModel_E3) = Flux.children(m)[1:7]
unpack_generative(m::MTSeqModels) = (m.x0decoder, m.mtbias_decoder, m.generator, m.dec_tform)

modelutils.load!(m::MTSeqModels, fname::String) = load!(Flux.params(m), fname)

function Base.show(io::IO, l::MTSeqModel_Pred)
    d_x0, d_mt = size(l.init_post.Dense1.W, 1), size(l.mt_post.Dense1.W, 1)
    out_type = Dict(MTSeq_CNN=>"CNN", MTSeq_Single=>"Deterministic", MTSeq_Double=>"Probabilistic")[model_type(l)]
    mttype = l.mtbias_only ? "Bias-only" : "Full-MT Recurrent Cell (" * get_strtype_wo_params(l.generator) * ")"
    print(io, "MTSeqModel_Pred(", mttype, ", d_x0=", d_x0, ", d_mt=", d_mt, ", ", out_type, ") -- for prediction")
end

function Base.show(io::IO, l::MTSeqModel_E3)
    d_x0, d_mt, d_c = size(l.init_post.Dense1.W, 1), size(l.mt_post.Dense1.W, 1), size(l.chaos_post.Dense1.W, 1)
    out_type = Dict(MTSeq_CNN=>"CNN", MTSeq_Single=>"Deterministic", MTSeq_Double=>"Probabilistic")[model_type(l)]
    mttype = l.mtbias_only ? "Bias-only" : "Full-MT Recurrent Cell (" * get_strtype_wo_params(l.generator) * ")"
    print(io, "MTSeqModel_E3(", mttype, ", d_x0=", d_x0, ", d_mt=", d_mt, ", d_chaos=", d_c, ", ", out_type, ") -- for llh")
end

################################################################################
##                                                                            ##
##                                Forward Model                               ##
##                                                                            ##
################################################################################

function (m::MTSeqModel_Pred)(y::AbstractVector; T_steps=70, T_enc=10, stoch=true)
    x0, z = posterior_samples(m, y; T_enc=T_enc, stoch=stoch)[1]
    return m(x0, z; T_steps=T_steps)
end

function (m::MTSeqModel_E3)(y::AbstractVector, yfull::AbstractVector; T_steps=70, T_enc=10, stoch=true)
    x0, z, c = posterior_samples(m, y, yfull; T_enc=T_enc, stoch=stoch)[1]
    return m(x0, z, c; T_steps=T_steps)
end

# ------------------- Forward Model (inference section) ----------------------

function posterior_samples(m::MTSeqModel_Pred, y::AbstractVector, yfull::AbstractVector=[];
    T_enc=10, stoch=true, x0=nothing, z=nothing, c=nothing)

    enc_tform, init_enc, init_post, mt_enc, mt_post = unpack_inference(m)
    (x0 === nothing || z === nothing) && (enc_yslices = [enc_tform(yy) for yy in y])

    # RNN 1: Amortized inference for initial state, h0, of generator
    if x0 === nothing
        x0, μ_x0, σ_x0 = _posterior_sample(init_enc, init_post, enc_yslices, T_enc, stoch)
    else
        μ_x0, σ_x0 = nothing, nothing
    end

    # RNN 2: Amortized inference for z
    # Technically this should be conditioned on the sample x₀, nevertheless, the posterior
    # of x0 is usually very tight, and the `mt_enc` has access to the same information.
    if z === nothing
        z, μ_z, σ_z = _posterior_sample(mt_enc, mt_post, enc_yslices, T_enc, stoch)
    else
        μ_z, σ_z = nothing, nothing
    end

    return (x0,z), (μ_x0, μ_z), (σ_x0, σ_z)
end

function posterior_samples(m::MTSeqModel_E3, y::AbstractVector, yfull::AbstractVector;
    T_enc=10, stoch=true, x0=nothing, z=nothing, c=nothing)

    enc_tform, init_enc, init_post, mt_enc, mt_post, chaos_enc, chaos_post = unpack_inference(m)

    # RNN 1: Amortized inference for initial state, h0, of generator
    if x0 === nothing
        enc_yslices = [enc_tform(yy) for yy in y]
        x0, μ_x0, σ_x0 = _posterior_sample(init_enc, init_post, enc_yslices, T_enc, stoch)
    else
        μ_x0, σ_x0 = nothing, nothing
    end

    # RNN 2: Amortized inference for z
    if z === nothing
        enc_fullseqs = [enc_tform(yy) for yy in yfull]
        z, μ_z, σ_z = _posterior_sample(mt_enc, mt_post, enc_fullseqs, length(yfull), stoch)
    else
        μ_z, σ_z = nothing, nothing
    end

    # RNN 3: Encode the chaos
    # Want a lowdim rep. of departure from expected trajected.
    if c === nothing
        !(x0 === nothing) && (enc_yslices = [enc_tform(yy) for yy in y])
        if !(z === nothing) && size(x0, 2) == 1 && size(z,2) > 1
            gmove = Flux.has_cuarrays() && Tracker.data(z) isa Flux.CuArray ? gpu : identity
            x0z = vcat(x0 * gmove(ones(Float32, 1, size(z,2))), z)
        else
            x0z = vcat(x0, z)
        end
        c, μ_c, σ_c = _posterior_sample(chaos_enc, chaos_post, enc_yslices, length(y),
            stoch, x0z)  # note conditioning on (x0, z) here.
    else
        μ_c, σ_c = nothing, nothing
    end
    return (x0,z,c), (μ_x0, μ_z, μ_c), (σ_x0, σ_z, σ_c)
end

# ------------------- Forward Model (generative section) ----------------------
# Note in Julia 1.1.1 (where this is developed), it is not possible to add
# methods to abstract types; hence (m::MTSeqModels)(x0, z, ...) is not possible
# and so we have the following "private"(!) function which is dispatched to by
# both types ∈ MTSeqModels. This is fixed in future versions of julia (I think)
# by https://github.com/JuliaLang/julia/pull/31916.
function _forward_given_latent(m::MTSeqModels, x0::AbstractVecOrMat, z::AbstractVecOrMat; T_steps=70)
    x0decoder, mtbias_decoder, gen_model, dec_tform = unpack_generative(m)

    if m.mtbias_only
        gen_model.state = x0decoder(x0)
        mtbias = mtbias_decoder(z)   # if linear, this is just the identity ∵ GRUCell.Wi
        return [dec_tform(gen_model(mtbias)) for t in 1:T_steps]
    else
        # Get GRU models from samples from the MT GRU model
        posterior_grus = gen_model(z) # output: BatchedGRU (def. above)
        # Set state (batch-wise) to x0 sample
        posterior_grus.state = x0decoder(x0)  # 2×linear d_x0 × nbatch → d_x × nbatch
        # Run generative model
        return [dec_tform(posterior_grus()) for t in 1:T_steps]
    end
end

(m::MTSeqModel_Pred)(x0::AbstractVecOrMat, z::AbstractVecOrMat; T_steps=70) =
    _forward_given_latent(m, x0, z; T_steps=T_steps)

# Allow signature for `E3` model to call with separate z, c vars..
(m::MTSeqModel_E3)(x0::AbstractVecOrMat, z::AbstractVecOrMat,
    c::AbstractVecOrMat; T_steps=70) = _forward_given_latent(m, x0, vcat(z, c); T_steps=T_steps)

# -------------------  Utils for inference section ----------------------------

function _posterior_sample(enc, dec, input, T_max, stochastic=true, input_cat=nothing)
    Flux.reset!(enc)
    if input_cat === nothing
        for tt = 1:T_max; enc(input[tt,:,:]); end
    else
        for tt = 1:T_max; enc(vcat(input[tt,:,:], input_cat)); end
    end
    μ_, σ_ = dec(enc.state)
    n, d = size(μ_)
    smp = μ_ + randn_repar(σ_, n, d, stochastic)
    return smp, μ_, σ_
end

function _posterior_sample(enc, dec, input::Vector, T_max, stochastic=true, input_cat=nothing)
    Flux.reset!(enc)
    if input_cat === nothing
        for tt = 1:T_max; enc(input[tt]); end
    else
        d, nbatch = size(input[1])
        if nbatch == 1 && nbatch != size(input_cat,2)
            gmove = Flux.has_cuarrays() && input_cat isa Flux.CuArray ? gpu : identity
            input_expander = gmove(ones(Float32, 1, size(input_cat, 2)))
        else
            input_expander = 1
        end
        for tt = 1:T_max; enc(vcat(input[tt] * input_expander, input_cat)); end
    end
    μ_, σ_ = dec(enc.state)
    n, d = size(μ_)
    smp = μ_ + randn_repar(σ_, n, d, stochastic)
    return smp, μ_, σ_
end

################################################################################
##                                                                            ##
##                             Model "Constructor"                            ##
##                                                                            ##
################################################################################

function create_model(d_x, d_x0, d_y, d_enc_state, d_mt, d_chaos=0; encoder=:GRU,
    cnn=false, out_heads=1, d_hidden=d_x, mtbias_only=false, d_hidden_mt=32,
    mt_is_linear=true, decoder_fudge_layer=false, model_purpose=:llh, fixb=false,
    fixb_version=nothing)

    @assert !(out_heads > 1 && cnn) "cannot have multiple output heads and CNN."
    @argcheck out_heads in 1:2
    @argcheck model_purpose in [:llh, :pred]
    @assert !(fixb && mtbias_only) "cannot do mtbias AND fix 'b'"
    @assert !fixb || something(fixb_version, 1) ∈ [1,2] "fixb_version must be in [1,2]."

    is_pred_model = model_purpose == :pred
    if fixb && fixb_version === nothing
        @warn "fixed bias chosen, but `fixb_version` not specified. Defaulting to v1."
        fixb_version = 1
    end

    # ENCODER TRANSFORM FROM OBS SPACE => ENC SPACE
    if cnn
        N_FILT=32
        d_conv_result = 4*4*N_FILT
        tform_enc = Chain(
            Conv((3, 3), 2=>N_FILT, relu, stride=2, pad=1),        # out: 16 × 16 × 32 × nbatch
            Conv((3, 3), N_FILT=>N_FILT, relu, stride=2, pad=1),   # out: 8 × 8 × 32 × nbatch
            Conv((3, 3), N_FILT=>N_FILT, relu, stride=2, pad=1),   # out: 4 × 4 × 32 × nbatch
            x->reshape(x, d_conv_result, :),
            Dense(d_conv_result, d_hidden, swish)
        )
        d_rnn_in = d_hidden
        model_type = MTSeq_CNN()
    else
        tform_enc = identity
        d_rnn_in = d_y
        model_type = out_heads == 1 ? MTSeq_Single() : MTSeq_Double()
    end

    # 2/3× ENCODERS for inference of
    # (a) initial state (init_enc)
    # (b) sequence level MT variable (mt_enc)
    # (c) local variability due to sensitivity/chaos (mt_chaos)
    @argcheck encoder in [:LSTM, :GRU, :Bidirectional]
    rnn_constructor = Dict(:LSTM=>LSTM, :GRU=>GRU, :Bidirectional=>BRNNenc)[encoder]
    init_enc = rnn_constructor(d_rnn_in, d_enc_state)             # (a)
    mt_enc = rnn_constructor(d_rnn_in, d_enc_state)               # (b)
    !(is_pred_model) && (chaos_enc = rnn_constructor(d_rnn_in+d_mt+d_x0, d_enc_state))  # (c)
    (encoder == :Bidirectional) && (d_enc_state *= 2)

    # POSTERIOR OVER LVM CORR. TO (a), (b), (c)
    init_post = MultiDense(Dense(d_enc_state, d_x0, identity), Dense(d_enc_state, d_x0, σ))
    init_post.Dense2.b.data .= -2   # initialize posteriors to be low variance
    mt_post = MultiDense(Dense(d_enc_state, d_mt, identity), Dense(d_enc_state, d_mt, σ))
    mt_post.Dense2.b.data .= -2
    if !is_pred_model
        chaos_post = MultiDense(Dense(d_enc_state, d_chaos, identity), Dense(d_enc_state, d_chaos, σ))
        chaos_post.Dense2.b.data .= -2
    end

    # decode from LV (a) --> size of generative hidden state
    x0decoder = Dense(d_x0, d_x, identity)

    # -------------- GENERATIVE MODEL -----------------
    ###################################################

    if !mtbias_only
        d_out = !fixb ? 3*d_x*(d_x+1) : (fixb_version == 1 ? 3*d_x^2 : 3*d_x^2 + 2*d_x)
        par_gen_net = mt_is_linear ? mlp(d_mt+d_chaos, d_out) : mlp(d_mt+d_chaos, d_hidden_mt, d_out; activation=tanh)
        if !fixb
            gen_rnn = MTGRU_NoU(d_x, par_gen_net)
        else
            if fixb_version == 1
                MTGRU_b = Flux.param(randn(Float32, d_x*3))
                gen_rnn = MTGRU_NoU_fixb(d_x, par_gen_net, MTGRU_b)
            elseif fixb_version == 2
                MTGRU_b = Flux.param(randn(Float32, d_x))
                gen_rnn = MTGRU_NoU_fixb_v2(d_x, par_gen_net, MTGRU_b)
            else
                error("fixb_version: Unreachable error.")
            end
        end
        mtbias_decoder = identity
    else
        mtbias_decoder = mt_is_linear ? identity : Dense(d_mt+d_chaos, d_hidden_mt, swish)
        _d_in = mt_is_linear ? d_mt+d_chaos : d_hidden_mt
        gen_rnn = GRU(_d_in, d_x)
    end

    if decoder_fudge_layer
        # accidental additional layer in decoder for full MT-model.
        # (This is really low capacity. 64->64 unit relu; perhaps as likely to reduce as to improve performance.)
        fudge_layer = Dense(d_x, d_x, relu)
    else
        fudge_layer = identity
    end

    if out_heads == 1 && !cnn
        decoder = Chain(Dense(d_x, d_hidden, relu), Dense(d_hidden, d_y, identity))
    elseif out_heads == 2 && !cnn
        decoder = Chain(Dense(d_x, d_hidden, relu), MultiDense(Dense(d_hidden, d_y, identity), Dense(d_hidden, d_y, identity)))
    else    # cnn
        decoder = Chain(
            fudge_layer,
            Dense(d_x, d_conv_result, identity),
            x->reshape(x, 4, 4, N_FILT, :),                               # out: 4 × 4 × 32 × nbatch
            ConvTranspose((3,3), N_FILT=>N_FILT, relu, stride=2),         # out: 9 × 9 × 16 × nbatch
            ConvTranspose((3,3), N_FILT=>N_FILT, relu, stride=2, pad=1),  # out: 17 × 17 × 8 × nbatch
            ConvTranspose((3,3), N_FILT=>2, identity, stride=2, pad=1),   # out: 33 × 33 × 1 × nbatch
            x->x[1:D_IM, 1:D_IM, :, :]   # out: 32 × 32 × 2 × nbatch
        )
    end

    if is_pred_model
        m = MTSeqModel_Pred(tform_enc, init_enc, init_post, mt_enc, mt_post,
            x0decoder, mtbias_decoder, gen_rnn, decoder, model_type, d_y, mtbias_only)
    else
        m = MTSeqModel_E3(tform_enc, init_enc, init_post, mt_enc, mt_post, chaos_enc, chaos_post,
            x0decoder, mtbias_decoder, gen_rnn, decoder, model_type, d_y, mtbias_only)
    end

    return m
end

"""
    create_model_opt_dict(...)
Takes same arguments as `create_model` but just wraps up arguments in a Dict for
saving out as JSON/YAML/BSON etc.
"""
function create_model_opt_dict(d_x, d_x0, d_y, d_enc_state, d_mt, d_chaos=0; encoder=:GRU,
    cnn=false, out_heads=1, d_hidden=d_x, mtbias_only=false, d_hidden_mt=32,
    mt_is_linear=true, decoder_fudge_layer=false, model_purpose=:llh, fixb=false,
    fixb_version=nothing)

    Dict("d_x"=>d_x, "d_x0"=>d_x0, "d_y"=>d_y, "d_enc_state"=>d_enc_state, "d_mt"=>d_mt, "d_chaos"=>d_chaos,
    "encoder"=>encoder, "cnn"=>cnn, "out_heads"=>out_heads, "d_hidden"=>d_hidden, "mtbias_only"=>mtbias_only,
    "d_hidden_mt"=>d_hidden_mt, "mt_is_linear"=>mt_is_linear, "decoder_fudge_layer"=>decoder_fudge_layer,
    "model_purpose"=>model_purpose, "fixb"=>fixb, "fixb_version"=>fixb_version)
end

"""
    load_model_from_def(ymlfile)
We have YAML files saved with metadata associated with the model parameters
in the /data folder. These YAML files point to the relevant parameters too, and
so suffices to construct the model and load the parameters. That is the purpose
of this function.
"""
function load_model_from_def(ymlfile::String)
    model_details = YAML.load_file(ymlfile)
    filename = model_details["filename"]
    constructor = model_details["constructor"]
    @argcheck constructor == "mtmodel.create_model"
    D = model_details["model_def"]
    D["encoder"] = Symbol(D["encoder"])
    D["model_purpose"] = Symbol(D["model_purpose"])
    m = create_model(D["d_x"], D["d_x0"], D["d_y"], D["d_enc_state"], D["d_mt"], D["d_chaos"];
        encoder=D["encoder"], cnn=D["cnn"], out_heads=D["out_heads"], d_hidden=D["d_hidden"],
        mtbias_only=D["mtbias_only"], d_hidden_mt=D["d_hidden_mt"], mt_is_linear=D["mt_is_linear"],
        decoder_fudge_layer=D["decoder_fudge_layer"], model_purpose=D["model_purpose"], fixb=D["fixb"],
        fixb_version=D["fixb_version"])
    modelutils.load!(m, filename)
    return m
end

# # Example saving model info
# YAML.write_file("saved_models/mtgru_video_fixbgrp_pred_450.yml", Dict(
#     "filename"=>"data/mtgru_video_fixbgrp_pred_450.bson",
#     "description"=>"MT GRU model with posterior mean GROUPED BY SEQ IDENTITY fixed (v2) bias. 450x2 epochs. T=80, Tenc=20.",
#     "constructor"=>"mtmodel.create_model",
#     "model_def"=>mtmodel.create_model_opt_dict(64, 6, 4, 40, 2; encoder=:GRU, cnn=true, d_hidden=128,
#     model_purpose=:pred, mtbias_only=false, decoder_fudge_layer=true, fixb=true, fixb_version=2)
# ))

################################################################################
##                                                                            ##
##                        Objective(s) / Loss Function(s)                     ##
##                                                                            ##
################################################################################

function online_inference_BCE(m::MTSeqModels, x0::AbstractVecOrMat, z::AbstractVecOrMat,
    y::AbstractVector; T_steps=length(y))

    x0decoder, mtbias_decoder, gen_model, dec_tform = unpack_generative(m)

    # Generative model
    # --------------------------------------------------
    if m.mtbias_only
        gen_model.state = x0decoder(x0)
        mtbias = mtbias_decoder(z)   # if linear, this is just the identity ∵ GRUCell.Wi

        nllh = map(1:T_steps) do t
            _nllh_bernoulli(dec_tform(gen_model(mtbias)), y[t])  # y broadcasts over the batch implicitly
        end
    else
        posterior_grus = gen_model(z) # output: BatchedGRU (def. above)
        posterior_grus.state = x0decoder(x0)  # 2×linear d_x0 × nbatch → d_x × nbatch

        nllh = map(1:T_steps) do t
            _nllh_bernoulli(dec_tform(posterior_grus()), y[t])  # y broadcasts over the batch implicitly
        end
    end
    return reduce(vcat, nllh)
end

# Allow signature for `E3` model to call with separate z, c vars..
online_inference_BCE(m::MTSeqModel_E3, x0::AbstractVecOrMat, z::AbstractVecOrMat, c::AbstractVecOrMat,
    y::AbstractVector; T_steps=length(y)) = online_inference_BCE(m, x0, vcat(z,c), y; T_steps=T_steps)

function online_inference_single_BCE(m::MTSeqModels, h0::AbstractVecOrMat, z::AbstractVecOrMat,
    y::AbstractArray)

    mtbias_decoder, gen_model, dec_tform = unpack(m)[9:11]

    # Generative model
    # --------------------------------------------------
    if m.mtbias_only
        gen_model.state = h0
        mtbias = mtbias_decoder(z)
        h_new = gen_model(mtbias)
        nllh = _nllh_bernoulli(dec_tform(h_new), y)
    else
        posterior_grus = gen_model(z) # output: BatchedGRU (def. above)
        posterior_grus.state = h0  # 2×linear d_x0 × nbatch → d_x × nbatch
        h_new = posterior_grus()
        nllh = _nllh_bernoulli(dec_tform(h_new), y)
    end
    return vec(nllh), h_new
end

# Allow signature for `E3` model to call with separate z, c vars..
online_inference_single_BCE(m::MTSeqModel_E3, h0::AbstractVecOrMat, z::AbstractVecOrMat,
    c::AbstractVecOrMat, y::AbstractArray) = online_inference_single_BCE(m, h0, vcat(z,c), y)

function _nllh_bernoulli(ŷ::AbstractVector{T}, y::AbstractVector{T}) where T <: AbstractArray{M,N} where {M,N}
    @assert N > 1 "Avoiding infinite recursion. " *
    "Current implementation does not permit nllh of bernoulli vector of vectors. Try unsqueezing."
    return [_nllh_bernoulli(ŷŷ, yy) for (ŷŷ, yy) in zip(ŷ, y)]
end

_nllh_bernoulli(ŷ::AbstractArray, y::AbstractArray) =
    let n=size(ŷ)[end]; res=Flux.logitbinarycrossentropy.(ŷ, y); sum(reshape(res, :, n), dims=1); end

_nllh_gaussian(ŷ_lσ::AbstractVector, y::AbstractVector) =
    [((ŷŷ, ll) = ŷl; δ=yy-ŷŷ; sum(δ.*δ./(exp.(2*ll)))/2+sum(ll)/2) for (ŷl, yy) in zip(ŷ_lσ, y)]

_nllh_gaussian_constvar(ŷ::AbstractVector, y::AbstractVector, logstd::Number) =
    [(let n=size(ŷ)[end]; δ=yy-ŷŷ; res=sum(δ.*δ./(exp.(2*logstd)))/2+sum(logstd)/2;
    sum(reshape(res, :, n), dims=1); end) for (ŷŷ, yy) in zip(ŷ, y)]

function nllh(m::MTSeqModels{U,A,B,V,W,S,N,M}, y::AbstractVector, yfull::AbstractVector=[];
    T_steps=70, T_enc=10, stoch=true) where {U,A,B,V,W,S,N,M <: MTSeq_CNN}
    x0_z_c = posterior_samples(m, y, yfull; T_enc=T_enc, stoch=stoch)[1]
    online_inference_BCE(m::MTSeqModels, x0_z_c..., y; T_steps=T_steps)
end

nllh(m::MTSeqModels{U,A,B,V,W,S,N,M}, y::AbstractVector, yfull::AbstractVector;
    T_steps=70, T_enc=10, stoch=true, logstd=-2.5) where {U,A,B,V,W,S,N,M <: MTSeq_Single} =
        vcat(_nllh_gaussian_constvar(m(y, yfull; T_steps=T_steps, T_enc=T_enc, stoch=stoch), y, logstd)...)

nllh(m::MTSeqModels{U,A,B,V,W,S,N,M}, y::AbstractVector, yfull::AbstractVector;
    T_steps=70, T_enc=10, stoch=true) where {U,A,B,V,W,S,N,M <: MTSeq_Double} =
        vcat(_nllh_gaussian(m(y, yfull; T_steps=T_steps, T_enc=T_enc, stoch=stoch), y)...)

nllh(::Type{MTSeq_CNN}, ŷ::AbstractVector, y::AbstractVector; logstd=-2.5) = _nllh_bernoulli(ŷ, y)
nllh(::Type{MTSeq_Single}, ŷ::AbstractVector, y::AbstractVector; logstd=-2.5) = _nllh_gaussian_constvar(ŷ, y, logstd)
nllh(::Type{MTSeq_Double}, ŷ::AbstractVector, y::AbstractVector; logstd=-2.5) = _nllh_gaussian(ŷ, y)

StatsBase.loglikelihood(m::MTSeqModels, y::AbstractVector, yfull::AbstractVector;
    T_steps=70, T_enc=10, stoch=true, logstd=-2.5) = -nllh(m, y, yfull, T_steps=T_steps,
        T_enc=T_enc, stoch=stoch, logstd=logstd)

function _gauss_kl_to_N01(μ::AbstractVecOrMat{T}, σ::AbstractVecOrMat{T}) where T
    T(0.5) * sum(1 .+ 2*log.(σ.*σ) - μ.*μ - σ.*σ)  # x.^2 -> x.*x due to old version of CuArrays / GPU (culiteral pow issue)
end
_gauss_kl_to_N01_deterministic(μ::AbstractVecOrMat{T}) where T  = T(0.5) * sum(1 .- μ.*μ)

function _kl_penalty_stoch(β::Vector{T}, μs::Vector{T}, σs::Vector{T}) where T <: Float32
    @argcheck length(β) == length(μs) == length(σs)
    kl = zero(T)
    for j in 1:length(β)
        kl += β[j] * _gauss_kl_to_N01(μs[j], σs[j])
    end
    return kl
end

function _kl_penalty_deterministic(β::Vector{T}, μs::Vector{T}) where T <: Float32
    @argcheck length(β) == length(μs)
    kl = zero(T)
    for j in 1:length(β)
        kl += β[j] * _gauss_kl_to_N01_deterministic(μs[j])
    end
    return kl
end

function elbo(m::MTSeqModels, y::AbstractVector, yfull::AbstractVector=[]; T_steps=70,
    T_enc=10, stoch=true, kl_coeff=1f0, βkl=ones(Float32, 3))

    smps, μs, σs = posterior_samples(m, y, yfull; T_steps=T_steps, T_enc=T_enc, stoch=stoch)
    model_out = m(x0, vcat(z, c); T_steps=T_steps)
    recon = -nllh(M, model_out, y)

    # ⇒ Initially no KL (Hard EM), and then an annealing sched. cf. Bowman et al. etc?
    kl = stoch ? _kl_penalty_stoch(β, μs, σs) : _kl_penalty_deterministic(β, μs)
    kl = kl * kl_coeff

    return - (recon + kl)
end

end

max_stars_repo_name: ornithos/mtds-dblpend

max_issues_repo_name: ornithos/mtds-dblpend

max_forks_repo_name: ornithos/mtds-dblpend


text:
## convert2array method:

"""
    convert2array(fld::gcmfaces)

Convert gcmfaces to array
"""
function convert2array(fld::gcmfaces)

if fld.grTopo=="llc";
    tmp1=cat(fld.f[1],fld.f[2],rotr90(fld.f[4]),rotr90(fld.f[5]);dims=1);
    tmp2=cat(rotl90(fld.f[3]),NaN*fld.f[3],NaN*fld.f[3],NaN*fld.f[3];dims=1);
    arr=cat(tmp1,tmp2;dims=2);
elseif fld.grTopo=="cs";
    tmp1=cat(fld.f[1],fld.f[2],rotr90(fld.f[4]),rotr90(fld.f[5]);dims=1);
    tmp2=cat(rotl90(fld.f[3]),NaN*fld.f[3],NaN*fld.f[3],NaN*fld.f[3];dims=1);
    tmp0=cat(NaN*fld.f[3],NaN*fld.f[3],NaN*fld.f[3],rotr90(fld.f[6]);dims=1);
    arr=cat(tmp0,tmp1,tmp2;dims=2);
elseif fld.grTopo=="ll";
  arr=fld.f[1];
else;
  error("unknown grTopo case");
end;

return arr;

end

## convert2array ##

"""
    convert2array(fld::Array{T,N}) where {T,N}

Convert array to gcmfaces
"""
function convert2array(fld::Array{T,N}) where {T,N}

grTopo=MeshArrays.grTopo;
nFaces=MeshArrays.nFaces;

v1=Array{Array{T,N},1}(undef,nFaces);
N>2 ? error("N>2 case not implemented yet") : nothing

if MeshArrays.grTopo=="llc";
    (n1,n2)=MeshArrays.facesSize[1];
    v1[1]=fld[1:n1,1:n2];
    v1[2]=fld[n1+1:n1*2,1:n2];
    v1[3]=rotr90(fld[1:n1,n2+1:n2+n1]);
    v1[4]=rotl90(fld[n1*2+1:n1*3,1:n2]);
    v1[5]=rotl90(fld[n1*3+1:n1*4,1:n2]);
elseif MeshArrays.grTopo=="cs";
    (n1,n2)=MeshArrays.facesSize[1];
    v1[1]=fld[1:n1,n1+1:n1+n2];
    v1[2]=fld[n1+1:2*n1,n1+1:n1+n2];
    v1[3]=rotr90(fld[1:n1,n1+n2+1:n2+n1*2]);
    v1[4]=rotl90(fld[n1*2+1:n1*3,n1+1:n1+n2]);
    v1[5]=rotl90(fld[n1*3+1:n1*4,n1+1:n1+n2]);
    v1[6]=rotl90(fld[n1*3+1:n1*4,1:n1]);
elseif MeshArrays.grTopo=="ll";
    v1[1]=fld;
else;
  error("unknown grTopo case");
end;

gcmfaces(nFaces,grTopo,v1);

end

## convert2gcmfaces method:

"""
    convert2gcmfaces(fld::gcmfaces)

Convert mitgcm output to gcmfaces
"""
function convert2gcmfaces(fld::gcmfaces)

    grTopo=MeshArrays.grTopo;
    nFaces=MeshArrays.nFaces;
    (n1,n2)=MeshArrays.ioSize;
    facesSize=MeshArrays.facesSize;

    aa=0;bb=0;
    for iFace=1:nFaces;
        aa=aa+prod(facesSize[iFace]);
        bb=bb+prod(size(fld.f[iFace]));
    end;
    n3=Int64(bb/aa);
    v11=NaN*zeros(aa,n3);

    i0=0; i1=0;
    for iFace=1:nFaces;
        i0=i1+1;
        nn=facesSize[iFace][1];
        mm=facesSize[iFace][2];
        i1=i1+nn*mm;
        v11[i0:i1,:]=reshape(fld.f[iFace],(nn*mm,n3));
    end;
    v1=reshape(v11,(n1,n2,n3));
    n3==1 ? v1=dropdims(v1,dims=3) : nothing

    return v1;

end

## convert2gcmfaces ##

"""
    convert2gcmfaces(fld::Array)

Convert gcmfaces to mitgcm output
"""
function convert2gcmfaces(fld::Array)

grTopo=MeshArrays.grTopo;
nFaces=MeshArrays.nFaces;
(n1,n2)=MeshArrays.ioSize;
facesSize=MeshArrays.facesSize;

n3=Int64(prod(size(fld))/n1/n2);

v0=reshape(fld,(n1*n2,n3));
v1=Array{Array{eltype(fld),ndims(fld)}}(undef,nFaces);
i0=0; i1=0;
for iFace=1:nFaces
  i0=i1+1;
  nn=facesSize[iFace][1]; mm=facesSize[iFace][2];
  i1=i1+nn*mm;
  if n3>1;
    v1[iFace]=reshape(v0[i0:i1,:],(nn,mm,n3));
  else;
    v1[iFace]=reshape(v0[i0:i1,:],(nn,mm));
  end;
end

gcmfaces(nFaces,grTopo,v1);

end


text:
using BinaryBuilder, Pkg

"""
    offset_version(upstream, offset)

Compute a version that allows distinguishing between changes in the upstream
version and changes to the JLL which retain the same upstream version.

When the `upstream` version is changed, `offset` version numbers should be reset
to `v"0.0.0"` and incremented following semantic versioning.
"""
function offset_version(upstream, offset)
    return VersionNumber(
        upstream.major * 100 + offset.major,
        upstream.minor * 100 + offset.minor,
        upstream.patch * 100 + offset.patch,
    )
end

# GCC version for building the whole system
gcc_version = v"6"

# Versions of various COIN-OR libraries

Bonmin_upstream_version = v"1.8.8"
Bonmin_gitsha = "65c56cea1e7c40acd9897a2667c11f91d845bb7b"
Bonmin_version_offset = v"0.0.0"
Bonmin_version = offset_version(Bonmin_upstream_version, Bonmin_version_offset)

Cbc_upstream_version = v"2.10.5"
Cbc_gitsha = "7b5ccc016f035f56614c8018b20d700978144e9f"
Cbc_version_offset = v"0.0.0"
Cbc_version = offset_version(Cbc_upstream_version, Cbc_version_offset)

Cgl_upstream_version = v"0.60.2"
Cgl_gitsha = "6377b88754fafacf24baac28bb27c0623cc14457"
Cgl_version_offset = v"0.0.0"
Cgl_version = offset_version(Cgl_upstream_version, Cgl_version_offset)

Clp_upstream_version = v"1.17.6"
Clp_gitsha = "756ddd3ed813eb1fa8b2d1b4fe813e6a4d7aa1eb"
Clp_version_offset = v"0.0.0"
Clp_version = offset_version(Clp_upstream_version, Clp_version_offset)

Osi_upstream_version = v"0.108.6"
Osi_gitsha = "dfa6449d6756fdd96912cf96e168d0be07b1d37c"
Osi_version_offset = v"0.0.0"
Osi_version = offset_version(Osi_upstream_version, Osi_version_offset)

CoinUtils_upstream_version = v"2.11.4"
CoinUtils_gitsha = "f709081c9b57cc2dd32579d804b30689ca789982"
CoinUtils_verson_offset = v"0.0.0"
CoinUtils_version =
    offset_version(CoinUtils_upstream_version, CoinUtils_verson_offset)

Ipopt_version = v"3.13.4"
Ipopt_gitsha = "3fbc0d29427d4290d15fb8842f78ddfbb929278f"

ALPS_upstream_version = v"1.5.7"
# This is not the exact 1.5.7 tag, but a few commits later on stable/1.5
ALPS_gitsha = "5b1a0b524979764d6ca929446762762712c035bb"
ALPS_version_offset = v"0.0.1"
ALPS_version = offset_version(ALPS_upstream_version, ALPS_version_offset)

BiCePS_version = v"0.9.4"
BiCePS_gitsha = "8e41545a3b1a36ca1d306b3af96f3804dc57e61f"

CHiPPS_BLIS_upstream_version = v"0.94.8"
# This is not the exact 0.94.8 tag, but a few commits later on releases/0.94.8
CHiPPS_BLIS_gitsha = "0ad6d4530d48e3b5f417a3407f5f2e8605145b50"
CHiPPS_BLIS_version_offset = v"0.0.1"
CHiPPS_BLIS_version =
    offset_version(CHiPPS_BLIS_upstream_version, CHiPPS_BLIS_version_offset)

SYMPHONY_version = v"5.6.17"
SYMPHONY_gitsha = "f917d42e6655a82ea4e9290aa7d41b0f60a91f20"

MibS_version = v"1.1.3"
MibS_gitsha = "4b7ec93c4bd1d6a978deff9987cf1df74f6598d3"

# Third-party packages needed by COIN-OR libraries.
ASL_version = v"0.1.2"
METIS_version = v"5.1.0"
MUMPS_seq_version = v"5.2.1"
OpenBLAS32_version = v"0.3.10"

# These are the platforms we will build for by default, unless further
# platforms are passed in on the command line
platforms = expand_cxxstring_abis(supported_platforms(;experimental=true))
platforms = filter!(!Sys.isfreebsd, platforms)

max_stars_repo_name: caleb-allen/Yggdrasil

max_issues_repo_name: caleb-allen/Yggdrasil

max_forks_repo_name: caleb-allen/Yggdrasil


text:
using Plots, ApproxFun; gr()

#undamped simple harmonic oscillator
#f =(x,y)-> [ y(2), -y(1) ]

#undamped pendulum
f =(x,y) -> [ y[2], -sin(y[1]) ]

#van der Pol oscillator. reaches coefficients length limit after a few iterations.
#μ = 1.
#f(x,y) = [ y[2], μ*(1-y[1]^2)*y[2] - y[1] ]

y0 = [3., 0.]
d = Interval(0.,20.)
x = Fun(identity, d)
y0=Fun(x->y0,d)
yn = y0
plot(yn[1])

while true
	err = sum([maximum(fun^2) for fun in vec(yn'-f(x,yn))])
	println("Length: $(length(yn)) \tMaximum point-wise squared error: $err")
	if err < 1e-10
		break
	end

    plot!(yn[1])
	integrand = Fun(f(x,yn))
    yn = chop(cumsum(integrand) + y0,100eps())
end

xts = d.a:0.01:d.b
plot(xts,yn[1])
plot!(xts,yn[2])

max_stars_repo_path: examples/Picard Iteration.jl

max_issues_repo_path: examples/Picard Iteration.jl

max_forks_repo_path: examples/Picard Iteration.jl


text:
# using KernelDensity
using Distributions
using myJuliaUtils
using Test

# univariate case

# reps=200
reps=20000000
x = randn(reps)
a=0:1
# @time y = kde(x)
# pdf(y, a)

@time z=fit(Normal, x)
a1=pdf.(z,a)

@test isapprox(a1,pdf.(Normal(0.,1.),a); rtol=1e-3)

q=3
w=4
e=reps
x=randn(q,w,e);
v=reshape(collect(1:1:q*w) ./10,q,w)
o= sum(v.-0.1 .<=x .<=v.+0.1,dims=3)./e ./0.2
# @time y=TVPVARPkg.useKDE2fit(x,v)
# @btime y=findKD($x,$v)

@time z=useN2fit(x,v)
@test isapprox(z, pdf.(Normal(0.,1.),v) ; rtol=1e-3)

# multivariate case
@time z=useMvN2fit(x,v)

f=similar(x,q)
for a=1:q
    f[a]=pdf(MvNormal(zeros(w),eye(w)), v[a,:])
end

@test isapprox(z, f ; rtol=1e-3)

max_stars_repo_name: zymbuzz/myJuliaUtils

max_issues_repo_name: zymbuzz/myJuliaUtils

max_forks_repo_name: zymbuzz/myJuliaUtils


text:
using SeisBeamforming
using Test

@testset "SeisBeamforming.jl" begin
    # Write your own tests here.
end

max_stars_repo_name: kura-okubo/SeisBeamforming

max_issues_repo_name: kura-okubo/SeisBeamforming

max_forks_repo_name: kura-okubo/SeisBeamforming


text:
# Copyright (c) 2015 Dahua Lin, Miles Lubin, Joey Huchette, Iain Dunning, and
#   contributors
#
# Use of this source code is governed by an MIT-style license that can be found
# in the LICENSE.md file or at https://opensource.org/licenses/MIT.

module Gurobi

import Pkg

# deps.jl file is always built via `Pkg.build`, even if we didn't find a local
# install and we want to use the artifact instead. This is so Gurobi.jl will be
# recompiled if we update the file. See issue #438 for more details.
include(joinpath(dirname(@__FILE__), "..", "deps", "deps.jl"))

if isdefined(@__MODULE__, :libgurobi)
    # deps.jl must define a local installation.
elseif Sys.islinux()
    # Let's use the artifact instead.
    const libgurobi = joinpath(
        Pkg.Artifacts.artifact"gurobilinux64",
        "gurobi951/linux64/lib/libgurobi95.so",
    )
else
    error("""
        Gurobi not properly installed. Please run Pkg.build(\"Gurobi\"). For
        more information go to https://github.com/jump-dev/Gurobi.jl
    """)
end

const _GUROBI_VERSION = if libgurobi == "__skipped_installation__"
    # The deps file is fake, with the intention to make Gurobi.jl loadable but
    # not usable.
    VersionNumber(9, 5, 1)
else
    let
        majorP, minorP, technicalP = Ref{Cint}(), Ref{Cint}(), Ref{Cint}()
        ccall(
            (:GRBversion, libgurobi),
            Cvoid,
            (Ptr{Cint}, Ptr{Cint}, Ptr{Cint}),
            majorP, minorP, technicalP,
        )
        VersionNumber(majorP[], minorP[], technicalP[])
    end
end

function _is_patch(x::VersionNumber, reference::VersionNumber)
    return x.major == reference.major && x.minor == reference.minor
end

if _is_patch(_GUROBI_VERSION, v"9.0")
    include("gen90/libgrb_common.jl")
    include("gen90/libgrb_api.jl")
elseif _is_patch(_GUROBI_VERSION, v"9.1")
    include("gen91/libgrb_common.jl")
    include("gen91/libgrb_api.jl")
elseif _is_patch(_GUROBI_VERSION, v"9.5")
    include("gen95/libgrb_common.jl")
    include("gen95/libgrb_api.jl")
else
    error("""
    You have installed version $_GUROBI_VERSION of Gurobi, which is not
    supported by Gurobi.jl. We require Gurobi version 9.0 or 9.1 or 9.5.

    After installing a supported version of Gurobi, run:

        import Pkg
        Pkg.rm("Gurobi")
        Pkg.add("Gurobi")

    Make sure you set the environment variable `GUROBI_HOME` following
    the instructions in the Gurobi.jl README, which is available at
    https://github.com/jump-dev/Gurobi.jl.

    If you have a newer version of Gurobi installed, changes may need to be made
    to the Julia code. Please open an issue at
    https://github.com/jump-dev/Gurobi.jl.
    """)
end

include("MOI_wrapper/MOI_wrapper.jl")
include("MOI_wrapper/MOI_callbacks.jl")
include("MOI_wrapper/MOI_multi_objective.jl")
include("MOI_wrapper/MOI_indicator_constraint.jl")

# Gurobi exports all `GRBXXX` symbols. If you don't want all of these symbols in
# your environment, then use `import Gurobi` instead of `using Gurobi`.

for sym in filter(s -> startswith("$s", "GRB"), names(@__MODULE__, all = true))
    @eval export $sym
end

end


text:
# ------------------------------------------------------------------------------------------
# ### Notebook 1
# #### 1.1
# Buscar `convert` y `parse` en la documentación.
# ------------------------------------------------------------------------------------------

?convert;
# ?parse

# ------------------------------------------------------------------------------------------
# #### 1.2
# Asignar `365` a una variable llamada `days`. Convierte `days` a flotante.
# ------------------------------------------------------------------------------------------

days = 365
convert(Float64, days)

# ------------------------------------------------------------------------------------------
# #### 1.3
# Ver que sucede ejectuando
#
# ```julia
# convert(Int64, '1')
# ```
# and
#
# ```julia
# parse(Int64, '1')
# ```
# ------------------------------------------------------------------------------------------

# Esto regresa el código ascii (un entero) asociado con el caracter '1'
convert(Int64, '1')

# Esto regresa un entero encapsulado entre comillas
parse(Int64, '1')

# ------------------------------------------------------------------------------------------
# ### Notebook 2
# #### 2.1
# Crea una cadena que dice "hola" 1000 veces
# ------------------------------------------------------------------------------------------

"hola"^1000

# ------------------------------------------------------------------------------------------
# #### 2.2
# Agrega dos números dentro de una cadena
# ------------------------------------------------------------------------------------------

m, n = 1, 1
"$m + $n = $(m + n)"

# ------------------------------------------------------------------------------------------
# ### Notebook 3
#
# #### 3.1
# Crea un arreglo, `arreglo`, que es un arreglo en 1D de 2 elementos, cada uno conteniendo
# el número 0.
# Agrega a `arreglo` para agregar un segundo número, `1`, a cada arreglo.
# ------------------------------------------------------------------------------------------

a_ray = [[0], [0]]
push!(a_ray[1], 1)
push!(a_ray[2], 1)
a_ray

# ------------------------------------------------------------------------------------------
# ### 3.2
# Trata de agregar "Emergencia" a `miagenda` con el valor `911`. Trata de agregar `911` como
# un entero y no como cadena. ¿Porqué no funciona?
# ------------------------------------------------------------------------------------------

miagenda = Dict("Jenny" => "867-5309", "Ghostbusters" => "555-2368")

miagenda["Emergency"] = 911
#=

Julia infiere que "miagenda" toma ambos llaves
y valores del tipo "String". Podemos ver que miagenda
es un Dict{String,String} con 2 entradas. Esto significa que
Julia no va a aceptar enteros como valores en miagenda.

=#

# ------------------------------------------------------------------------------------------
# #### 3.3
# Crea un nuevo diccionario que se llame `agenda_flexible` que tenga el número de Jenny
# guardado como cadena y el de los Cazafantasmas como entero.
# ------------------------------------------------------------------------------------------

agenda_flexible = Dict("Jenny" => "867-5309", "Ghostbusters" => 5552368)

# ------------------------------------------------------------------------------------------
# #### 3.4
# Add the key "Emergency" with the value `911` (an integer) to `flexible_phonebook`.
# ------------------------------------------------------------------------------------------

flexible_phonebook["Emergency"] = 911

# ------------------------------------------------------------------------------------------
# ##### 3.5
# 3.5 ¿Porqué podemos agregar un entero como valor a `agenda_flexible` pero no a
# `miagenda`? ¿Cómo pudimos haber inicializado `miagenda` para que aceptara enteros como
# valores?
# ------------------------------------------------------------------------------------------

#=

Julia infiere que miagenda_flexible toma valores del tipo
Any. A diferencia de miagenda, miagenda_flexible es un
Dict{String,Any} con 2 entradas.

Para evitar esto, podemos inicializar miagenda a un
diccionario vacío y agregamos entradas después. O podemos
decirle a Julia explícitamente que queremos un diccionario
que acepte objectos del tipo Any como valores.

¡Ve los ejemplos!

=#

miagenda = Dict()

miagenda = Dict{String, Any}("Jenny" => "867-5309", "Ghostbusters" => "555-2368")

# ------------------------------------------------------------------------------------------
# ### Notebook 4
#
# #### 4.1
#
# 4.1 Crea un diccionario `squares`, que tiene llaves de valores de 1 a 100. El valor
# asociado a cada llave es el cuadrado de la llave. Guarda los valores asociados a las
# llaves pares como enteros y las impares como cadenas. Por ejemplo,
#
# ```julia
# squares[10] == 100
# squares[11] == "121"
# ```
#
# (¡No necesitas condicionales para esto!)
# ------------------------------------------------------------------------------------------

squares = Dict()
iterable = range(1, 2, 50)
#  Otra opción
# iterable = 1:2:99
for key in iterable
    squares[key] = "$(key^2)"
    squares[key + 1] = (key + 1)^2
end

@show squares[10]
@show squares[11]

# ------------------------------------------------------------------------------------------
# #### 4.2
#
# 4.2 Usa `fill` para crea una matriz de `10x10` de solo `0`'s. Pobla las primeras 10
# entradas con el índice de esa entrada. ¿Julia usa el orden de primero columna o primero
# renglón? (O sea, ¿el "segundo" elemento es el de la primera columna en el primer renglón,
# ó es el de el primer renglón en la segunda columna?)
# ------------------------------------------------------------------------------------------

A = fill(0, (10, 10))
for i in 1:10
    A[i] = i
end
A
# ¡Julia usa order columnar! Los primeros 10 elementos van a poblar la primera columna.

# ------------------------------------------------------------------------------------------
# ### Notebook 5
#
# #### 5.1
#
# 5.1 Reescribe FizzBuzz sin usar `elseif`.
# ------------------------------------------------------------------------------------------

N = 16
if (N % 3 == 0) & (N % 5 == 0)
    println("FizzBuzz")
else
    if (N % 3 == 0)
        println("Fizz")
    else
        if (N % 5 == 0)
            println("Buzz")
        else
            println(N)
        end
    end
end

# ------------------------------------------------------------------------------------------
# #### 5.2
#
# Reescribe FizzBuzz usando el operador ternario.
# ------------------------------------------------------------------------------------------

((N % 3 == 0) & (N % 5 == 0)) ? println("FizzBuzz") : ((N % 3 == 0) ? println("Fizz") : ((N % 5 == 0) ? println("Buzz") : println(N)))

# ------------------------------------------------------------------------------------------
# ### Notebook 6
#
# #### 6.1
#
# 6.1 En vez de broadcastear `f` sobre `v`,  pudimos haber hecho `v .^ 2`.
#
# Sin declarar una nueva funcion, agrega 1 a cada elemento de una matriz de `3x3` llena de
# `0`'s.
# ------------------------------------------------------------------------------------------

fill(0, (3, 3)) .+ 1

# ------------------------------------------------------------------------------------------
# #### 6.2
#
# 6.2 En vez de broadcastear `f` sobre el vector `v` con la sintaxis de punto, aplica `f` a
# todos los elementos de`v` usando `map` como función.
# ------------------------------------------------------------------------------------------

f(x) = x^2
v = [1, 2, 3]
map(f, v)

# ------------------------------------------------------------------------------------------
# #### 6.3
#
# Una cifra de César recorre cada letra un número determinado de plazas más adelante en el
# abecedario. Un corrimiento, o shift, de 1 manda "A" a "B". Escribe una función llamada
# `cesar` que toma una cadena como input y un corrimiento y regresa una cadena desencriptada
# tal que obtengas
#
# ```julia
# cesar("abc", 1)
# "bcd"
#
# cesar("hello", 4)
# "lipps"
# ```
# ------------------------------------------------------------------------------------------

caesar(input_string, shift) = map(x -> x + shift, input_string)

# ------------------------------------------------------------------------------------------
# ### Notebook 7
#
# #### 7.1
#
# 7.1 Usa el paquete de  (código fuente en https://github.com/JuliaMath/Primes.jl) para
# encontrar el número primer más grande menor a 1,000,000
# ------------------------------------------------------------------------------------------

#Pkg.add("Primes")
using Primes
maximum(primes(1000000))

# ------------------------------------------------------------------------------------------
# ### Notebook 8
#
# #### 8.1
#
# 8.1 Grafica y vs x para `y = x^2` usando el backend de PyPlot.
# ------------------------------------------------------------------------------------------

using Plots
pyplot()
x = 1:10
y = x .^ 2
plot(x, y)

# ------------------------------------------------------------------------------------------
# ### Notebook 9
#
# #### 9.1
#
# Agrega un método para `+` que aplique un cifrado de César a una cadena (cómo en el
# notebook 6) tal que
#
# ```julia
# "hello" + 4 == "lipps"
# ```
# ------------------------------------------------------------------------------------------

import Base: +
+(x::String, y::Int) = map(x -> x + y, x)

# ------------------------------------------------------------------------------------------
# #### 9.2
#
# Checa que has extendido propiamente `+` recorriendo la próxima cadena para atrás por 3
# letras:
#
# "Gr#qrw#phggoh#lq#wkh#diidluv#ri#gudjrqv#iru#|rx#duh#fuxqfk|#dqg#wdvwh#jrrg#zlwk#nhwfkxs1"
# ------------------------------------------------------------------------------------------

"Gr#qrw#phggoh#lq#wkh#diidluv#ri#gudjrqv#iru#|rx#duh#fuxqfk|#dqg#wdvwh#jrrg#zlwk#nhwfkxs1" + -3

# ------------------------------------------------------------------------------------------
# ### Notebook 10
#
# #### 10.1
#
# 10.1 Usa `circshift` para obtener una matriz con las columnas de A cíclicamente recorridas
# a la derecha por 3 columnas.
#
# Empezando con
#
# ```
# A = [
#  1  2  3  4  5  6  7  8  9  10
#  1  2  3  4  5  6  7  8  9  10
#  1  2  3  4  5  6  7  8  9  10
#  1  2  3  4  5  6  7  8  9  10
#  1  2  3  4  5  6  7  8  9  10
#  1  2  3  4  5  6  7  8  9  10
#  1  2  3  4  5  6  7  8  9  10
#  1  2  3  4  5  6  7  8  9  10
#  1  2  3  4  5  6  7  8  9  10
#  1  2  3  4  5  6  7  8  9  10
#  ]
# ```
#
# Quieres obtener
#
# ```
# A = [
#  7  8  9  10  1  2  3  4  5  6
#  7  8  9  10  1  2  3  4  5  6
#  7  8  9  10  1  2  3  4  5  6
#  7  8  9  10  1  2  3  4  5  6
#  7  8  9  10  1  2  3  4  5  6
#  7  8  9  10  1  2  3  4  5  6
#  7  8  9  10  1  2  3  4  5  6
#  7  8  9  10  1  2  3  4  5  6
#  7  8  9  10  1  2  3  4  5  6
#  7  8  9  10  1  2  3  4  5  6
#  ]
# ```
# ------------------------------------------------------------------------------------------

circshift(A, (0, 4))

# ------------------------------------------------------------------------------------------
# #### 10.2
#
# 10.2 Toma el producto de un vector `v` con sí mismo.
# ------------------------------------------------------------------------------------------

v = [1, 2, 3]
v * v'

# ------------------------------------------------------------------------------------------
# #### 10.3
# 10.3 Toma el producto de un vector `v` con sí mismo.
# ------------------------------------------------------------------------------------------

v' * v

# ------------------------------------------------------------------------------------------
# ### Notebook 11
#
# #### 11.1
#
#  ¿Cuáles son los eigenvalores de la Matriz A
#
# ```
# A =
# [
#  140   97   74  168  131
#   97  106   89  131   36
#   74   89  152  144   71
#  168  131  144   54  142
#  131   36   71  142   36
# ]
# ```
# ------------------------------------------------------------------------------------------

A =
[
 140   97   74  168  131
  97  106   89  131   36
  74   89  152  144   71
 168  131  144   54  142
 131   36   71  142   36
]

eigdec = eigfact(A)
eigdec[:values]

# ------------------------------------------------------------------------------------------
# #### 11.2
#
# Crea una matriz diagonal de los eigenvalores de A
# ------------------------------------------------------------------------------------------

Diagonal(eigdec[:values])

# ------------------------------------------------------------------------------------------
# #### 11.3
#
# Realiza un factorización de Hessenberg sobre la matriz A. Verifica que `A = QHQ'`.
# ------------------------------------------------------------------------------------------

F = hessfact(A)

isapprox(A, F[:Q] * F[:H] * F[:Q]')

max_stars_repo_name: grenkoca/JuliaTutorials

max_issues_repo_name: grenkoca/JuliaTutorials

max_forks_repo_name: grenkoca/JuliaTutorials


text:
import Base: A_mul_Bc!, A_mul_B!, (*)
import Base: getindex, setindex!

abstract type Dummy end
AbstractM  = Union{AbstractMatrix, LinearAlgebra.AbstractRotation, Dummy}
BigFloatOrComplex = Union{Complex{BigFloat}, BigFloat}

#=
"""
Add missing multiplication to LinearAlgebra.Rotation
"""
function A_mul_Bc!(R::LinearAlgebra.Rotation, G::LinearAlgebra.Givens)
  insert!(R.rotations, 1, G')
  R
end

function A_mul_B!{T<:Number,S<:Number}(A::AbstractMatrix{T}, R::LinearAlgebra.Rotation{S})
  n = length(R.rotations)
  @inbounds for i = 1:n
    A_mul_Bc!(A, R.rotations[n+1-i]')
  end
  A
end

function A_mul_B!(R::LinearAlgebra.Rotation, A::AbstractVecOrMat)
    @inbounds for i = 1:length(R.rotations)
        A_mul_B!(R.rotations[i], A)
    end
    return A
end

(*){T<:Number,S<:Number}(A::AbstractMatrix{T}, R::LinearAlgebra.Rotation{S}) = A_mul_B!(copy(A), R)

A_mul_Bc!(Q::Dummy, ::Any) = Q
A_mul_B!(Q::Dummy, ::Any) = Q
A_mul_B!(A::Any, Q::Dummy) = Q
(*)(Q::Dummy, A::Any) = Q
(*)(A::Any, Q::Dummy) = Q

getindex(Q::Dummy, ::Any, ::Any) = Q
setindex!(Q::Dummy, ::Any, ::Any, ::Any) = Q
=#

"""
Discriminant of 2 x 2 submatrix
"""
@inline discriminant(A, i, j) = ((A[i,i] - A[j,j]) / 2 ) ^2 + A[i,j] * A[j,i]

"""
Extract eigenvalues from quasi triangular matrix
"""
function seigendiag(A::AbstractMatrix{T}, ilo::Integer, ihi::Integer) where {T<:Real}
  # ev = eig(A[ilo:ihi,ilo:ihi])[1]  # TODO replace by specialized eig
  # filter( x -> imag(x) >= 0, ev)
  ev = Complex{eltype(A)}[]
  k = ilo
  while k <= ihi
    if k >= ihi || A[k+1,k] == 0
      push!(ev, A[k,k])
    else
      r = ( A[k,k] + A[k+1,k+1] ) / 2
      disc = discriminant(A, k, k + 1)
      if disc < 0
        dd = sqrt(-disc) * im
        push!(ev, r + dd)
        push!(ev, r - dd)
      else
        dd = sqrt(disc)
        push!(ev, r + dd)
        push!(ev, r - dd)
      end
      k += 1
    end
    k += 1
  end
  isreal(ev) ? real(ev) : ev
end
function seigendiag(A::AbstractMatrix{T}, ilo::Integer, ihi::Integer) where {T<:Complex}
  ev = diag(A[ilo:ihi,ilo:ihi])
  isreal(ev) ? real(ev) : ev
end

"""
Calculate vector of eigenvalues of a real 1x1 or 2x2 matrix A.
If eigenvalue is complex, return one complex number with positive imaginary part.
"""
function eig2(A::AbstractMatrix)
  n = size(A, 1)
  if n == 1
    [ A[1,1] ]
  else
    disc = discriminant(A, 1, 2)
    re = (A[1,1] + A[2,2]) / 2
    sq = sqrt(abs(disc))
    if disc < 0
      [ Complex(re, sq) ]
    else
      [ re - sq; re + sq ]
    end
  end
end

"""

  `is_hessenberg(::AbstractMatrix) -> bool`

  Test if matrix is exact Hessenberg form.
  No square matrix required.
  (Below subdiagonal exact zero except optional spike column)
"""
function is_hessenberg(A::AbstractMatrix, ispike::Integer = 0)
  n, m = size(A)
  for k = 1:m
    if k != ispike
      for i = k+2:n
        if A[i,k] != 0
          return false
        end
      end
    end
  end
  true
end

"""
Finish up to standardized 2x2 diagonal blocks
"""
function finish!(A::AbstractMatrix, Q::AbstractM)
  n = size(A, 1)
  for k = 1:n-1
    if A[k+1,k] != 0
      r = k:k+1
      AA, G = givens1(A, k, k+1)
      A_mul_B!(G, view(A, :, k+1:n))
      A_mul_Bc!(view(A, 1:k-1, :), G)
      A_mul_Bc!(Q, G)
      A[r,r] = AA
      # println("finish($k:$(k+1))")
      # display(@view A[k:k+1,k:k+1])
      # display(AA)
    end
  end
end

"""
Set elements of matrix to zero, which should be, but are not due to numerical errors.
"""
@inline function reschur!(A::AbstractMatrix, k::Integer)
  A[k+1:end,1:k] = 0
end

"""
Produce Givens Rotation, which transforms 2x2 matrix to another 2x2 matrix with
either: subdiagonal zero if 2 real eigenvalues.
The eigenvalue of lowest absolute value is placed in lower right position.
or: equal diagonal elements if 2 non-real complex eigenvalues.
The lower left element is <= upper right element absolutely.
"""
function givens1(A::AbstractMatrix{T}, i1::Integer, i2::Integer) where {T<:Number}
  a, b, x, d = A[i1,i1], A[i1,i2], A[i2,i1], A[i2,i2]
  btx = b * x
  apd = ( a + d ) / 2
  bpx = ( b + x ) / 2
  da  = ( d - a ) / 2
  disc = da ^ 2 + btx
  if disc >= 0
    root = sqrt( disc )
    root = copysign(root, apd)
    e1 = apd + root
    e2 = ( a * d - btx ) / e1
    ea = a - e2
    ed = d - e2
    G, r = abs(b) + abs(ed) > abs(b) + abs(ea) ? givens(b, ed, i1, i2) : givens(ea, x, i1, i2)
    v = 2( da * G.c - bpx * G.s) * G.s + b
    [e1 v; 0 e2], G
  else
    bx = ( b - x ) / 2
    root = hypot(da, bpx)
    root = copysign(root, bx)
    w = bx + root
    v = disc / w
    G, r = abs(b) >= abs(x) ? givens(w + x, da, i1, i2) : givens(da, w - b, i1, i2)
    [ apd w; v apd], G
  end
end

max_stars_repo_name: KlausC/LinAlgBigFloat

max_issues_repo_name: KlausC/LinAlgBigFloat

max_forks_repo_name: KlausC/LinAlgBigFloat


text:
using Documenter, FeedbackNets, Flux

makedocs(
    sitename = "FeedbackNets Documentation",
    modules = [FeedbackNets, Flux],
    pages = [
        "Home" => "index.md"
        "Guide" => [
            "Getting Started" => "guide/getting_started.md",
            "Chains vs Trees" => "guide/chains_vs_trees.md",
            "Working with Networks" => "guide/working_with_networks.md"
        ]
        "Reference" => "reference.md"
    ]
)

deploydocs(
    repo = "github.com/cJarvers/FeedbackNets.jl.git"
)


text:
using Agents
using Statistics: mean
using Random # hide

mutable struct Daisy <: AbstractAgent
    id::Int
    pos::Dims{2}
    breed::Symbol
    age::Int
    albedo::Float64 # 0-1 fraction
end

const DaisyWorld = ABM{<:GridSpace, Daisy};

function update_surface_temperature!(pos, model::DaisyWorld)
    ids = ids_in_position(pos, model)
    absorbed_luminosity = if isempty(ids) # no daisy
        ## Set luminosity via surface albedo
        (1 - model.surface_albedo) * model.solar_luminosity
    else
        ## Set luminosity via daisy albedo
        (1 - model[ids[1]].albedo) * model.solar_luminosity
    end
    ## We expect local heating to be 80 ᵒC for an absorbed luminosity of 1,
    ## approximately 30 for 0.5 and approximately -273 for 0.01.
    local_heating = absorbed_luminosity > 0 ? 72 * log(absorbed_luminosity) + 80 : 80
    ## Surface temperature is the average of the current temperature and local heating.
    model.temperature[pos...] = (model.temperature[pos...] + local_heating) / 2
end

function diffuse_temperature!(pos, model::DaisyWorld)
    ratio = get(model.properties, :ratio, 0.5) # diffusion ratio
    npos = nearby_positions(pos, model)
    model.temperature[pos...] =
        (1 - ratio) * model.temperature[pos...] +
        ## Each neighbor is giving up 1/8 of the diffused
        ## amount to each of *its* neighbors
        sum(model.temperature[p...] for p in npos) * 0.125 * ratio
end

function propagate!(pos, model::DaisyWorld)
    ids = ids_in_position(pos, model)
    if !isempty(ids)
        daisy = model[ids[1]]
        temperature = model.temperature[pos...]
        ## Set optimum growth rate to 22.5 ᵒC, with bounds of [5, 40]
        seed_threshold = (0.1457 * temperature - 0.0032 * temperature^2) - 0.6443
        if rand(model.rng) < seed_threshold
            ## Collect all adjacent position that have no daisies
            empty_neighbors = Tuple{Int,Int}[]
            neighbors = nearby_positions(pos, model)
            for n in neighbors
                if isempty(ids_in_position(n, model))
                    push!(empty_neighbors, n)
                end
            end
            if !isempty(empty_neighbors)
                ## Seed a new daisy in one of those position
                seeding_place = rand(model.rng, empty_neighbors)
                add_agent!(seeding_place, model, daisy.breed, 0, daisy.albedo)
            end
        end
    end
end

function daisy_step!(agent::Daisy, model::DaisyWorld)
    agent.age += 1
    agent.age >= model.max_age && kill_agent!(agent, model)
end

function daisyworld_step!(model)
    for p in positions(model)
        update_surface_temperature!(p, model)
        diffuse_temperature!(p, model)
        propagate!(p, model)
    end
    model.tick += 1
    solar_activity!(model)
end

function solar_activity!(model::DaisyWorld)
    if model.scenario == :ramp
        if model.tick > 200 && model.tick <= 400
            model.solar_luminosity += model.solar_change
        end
        if model.tick > 500 && model.tick <= 750
            model.solar_luminosity -= model.solar_change / 2
        end
    elseif model.scenario == :change
        model.solar_luminosity += model.solar_change
    end
end

import StatsBase
import DrWatson: @dict
using Random

function daisyworld(;
    griddims = (30, 30),
    max_age = 25,
    init_white = 0.2, # % cover of the world surface of white breed
    init_black = 0.2, # % cover of the world surface of black breed
    albedo_white = 0.75,
    albedo_black = 0.25,
    surface_albedo = 0.4,
    solar_change = 0.005,
    solar_luminosity = 1.0, # initial luminosity
    scenario = :default,
    seed = 165,
)

    rng = MersenneTwister(seed)
    space = GridSpace(griddims)
    properties = @dict max_age surface_albedo solar_luminosity solar_change scenario
    properties[:tick] = 0
    properties[:temperature] = zeros(griddims)

    model = ABM(Daisy, space; properties, rng)

    ## Populate with daisies: each position has only one daisy (black or white)
    grid = collect(positions(model))
    num_positions = prod(griddims)
    white_positions =
        StatsBase.sample(grid, Int(init_white * num_positions); replace = false)
    for wp in white_positions
        wd = Daisy(nextid(model), wp, :white, rand(model.rng, 0:max_age), albedo_white)
        add_agent_pos!(wd, model)
    end
    allowed = setdiff(grid, white_positions)
    black_positions =
        StatsBase.sample(allowed, Int(init_black * num_positions); replace = false)
    for bp in black_positions
        wd = Daisy(nextid(model), bp, :black, rand(model.rng, 0:max_age), albedo_black)
        add_agent_pos!(wd, model)
    end

    ## Adjust temperature to initial daisy distribution
    for p in positions(model)
        update_surface_temperature!(p, model)
    end

    return model
end

# ## Visualizing & animating
# %% #src

using InteractiveDynamics
using GLMakie

model = daisyworld()

daisycolor(a::Daisy) = a.breed

plotkwargs = (
    ac=daisycolor, as = 12, am = '♠',
    heatarray = :temperature,
    heatkwargs = (colorrange = (-20, 60),),
)
fig, _ = abm_plot(model; plotkwargs...)
fig

# And after a couple of steps
Agents.step!(model, daisy_step!, daisyworld_step!, 5)
fig, _ = abm_plot(model; heatarray = model.temperature, plotkwargs...)
fig

# %% Video
model = daisyworld()
abm_video(
    "daisyworld.mp4",
    model,
    daisy_step!,
    daisyworld_step!;
    title = "Daisy World",
    plotkwargs...,
)

# ## Interactive
# %% #src
using InteractiveDynamics, GLMakie, Random
model = daisyworld(; solar_luminosity = 1.0, solar_change = 0.0, scenario = :change)

black(a) = a.breed == :black
white(a) = a.breed == :white
adata = [(black, count), (white, count)]
temperature(model) = mean(model.temperature)
mdata = [temperature, :solar_luminosity]

params = Dict(
    :surface_albedo => 0:0.01:1,
    :solar_change => -0.1:0.01:0.1,
)
alabels = ["black", "white"]
mlabels = ["T", "L"]

fig, adf, mdf = abm_data_exploration(
    model, daisy_step!, daisyworld_step!, params;
    mdata, adata, alabels, mlabels, plotkwargs...
)

display(fig)


text:
# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

module GeoStatsBase

using CSV
using Meshes
using Tables
using TableOperations
using LinearAlgebra: normalize
using Combinatorics: multiexponents
using Distributions: median, mode
using Distributed: CachingPool, pmap, myid
using Distances: Euclidean, evaluate, pairwise
using StatsBase: Histogram, AbstractWeights, midpoints, sample
using Transducers: Map, foldxt
using StaticArrays: SVector
using DensityRatioEstimation
using ScientificTypes
using LossFunctions
using RecipesBase
using Parameters
using ReferenceFrameRotations

using TypedTables # for a default table type
using Optim # for LSIF estimation

import Meshes
import MLJModelInterface
import StatsBase: fit, varcorrection
import Statistics: mean, var, quantile

# aliases
const MI = MLJModelInterface
const GeoData = Meshes.MeshData

include("georef.jl")
include("ensembles.jl")
include("macros.jl")
include("trends.jl")
include("estimators.jl")
include("weighting.jl")
include("geoops.jl")
include("learning.jl")
include("mappings.jl")
include("problems.jl")
include("solvers.jl")
include("folding.jl")
include("errors.jl")
include("statistics.jl")
include("histograms.jl")
include("rotations.jl")
include("plotrecipes.jl")

export
  # data
  GeoData,
  georef,

  # ensembles
  Ensemble,
  domain,

  # mapping
  MappingMethod,
  NearestMapping,
  CopyMapping,

  # learning tasks
  LearningTask,
  SupervisedLearningTask,
  UnsupervisedLearningTask,
  RegressionTask,
  ClassificationTask,
  ClusteringTask,
  issupervised,
  inputvars,
  outputvars,
  features,
  label,

  # learning models
  issupervised,
  isprobabilistic,
  iscompatible,

  # learning losses
  defaultloss,

  # problems
  Problem,
  EstimationProblem,
  SimulationProblem,
  LearningProblem,
  data, domain,
  sourcedata,
  targetdata,
  task,
  variables,
  coordinates,
  hasdata,
  nreals,

  # solvers
  Solver,
  EstimationSolver,
  SimulationSolver,
  LearningSolver,
  targets,
  covariables,
  preprocess,
  solve, solvesingle,

  # folding
  FoldingMethod,
  RandomFolding,
  PointFolding,
  BlockFolding,
  BallFolding,
  folds,

  # errors
  ErrorEstimationMethod,
  LeaveOneOut,
  LeaveBallOut,
  CrossValidation,
  BlockCrossValidation,
  WeightedCrossValidation,
  DensityRatioValidation,

  # helper macros
  @estimsolver,
  @simsolver,

  # estimators
  fit, predict, status,

  # weighting
  GeoWeights,
  WeightingMethod,
  UniformWeighting,
  BlockWeighting,
  DensityRatioWeighting,
  weight,

  # trends
  polymat,
  trend,

  # histograms
  EmpiricalHistogram,

  # statistics
  mean, var, quantile,

  # plot recipes
  cornerplot,

  # utilities
  uniquecoords,
  groupby,

  # rotations
  DatamineAngles,
  GslibAngles,
  VulcanAngles

end


text:
# https://github.com/FluxML/Flux.jl/issues/136
include("ADVI.jl")

using Flux, Flux.Tracker
using Distributions

x = param([.1])
p = ADVI.SB_transform(x)
ADVI.SB_logabsdetJ(x, p)

x = param([.1, .2, .3])
p = ADVI.SB_transform(x)
@assert abs(ADVI.SB_logabsdetJ(x, p)[1] - (-5.5850)) < 1e-4
@assert abs(logpdf(Dirichlet(ones(4)), p) + ADVI.SB_logabsdetJ(x, p)[1] - (-3.7933)) < 1e-4

# TEST
y = param(ones(Float32, 3))
p = ADVI.SB_transform(y)
z = sum(log.(p))
back!(z)
@assert all(abs.(y.tracker.grad - [-0.9015, -0.7284, -0.4621]) .< 1e-4)
# 3-element Array{Float64,1}:
# -0.9014675456746869
# -0.7283506542974874
# -0.4621171572600099

y = param(ones(Float32, 3))
p = ADVI.SB_transform(y)
z = ADVI.logsumexp(p)[1]
back!(z)
y.tracker.grad

x = cumsum(cumsum(cumsum(ones(2,3,4), dims=1), dims=2), dims=3)
x = param(x / 24)
p = ADVI.SB_transform(x)
z = sum(log.(p))
back!(z)
x.grad
@assert abs(sum(x.grad) - (-4.8753)) < 1e-4

x = cumsum(cumsum(cumsum(ones(2,3,4), dims=1), dims=2), dims=3)
x = param(x / 24)
p = ADVI.SB_transform(x)
z = ADVI.logsumexp(p, dims=3)
back!(sum(z))
x.grad

# These were tested against pytorch
# ADVI.SB_logabsdetJ(x, p)
# ADVI.compute_lpdf(Dirichlet(collect(1:5)), p)

println("Done.")


text:
using GLPK

JuDGE_SP_Solver = optimizer_with_attributes(GLPK.Optimizer, "msg_lev" => 0, "mip_gap" => 0.0)
JuDGE_MP_Solver = optimizer_with_attributes((method=GLPK.INTERIOR) -> GLPK.Optimizer(), "msg_lev" => 0, "mip_gap" => 0.0)
JuDGE_DE_Solver = optimizer_with_attributes(GLPK.Optimizer, "msg_lev" => 2, "mip_gap" => 0.0)


text:
include("Graph.jl")

function generatePicture(X, Y, labels, PATH, PATH2, NodeC, na, maxd)
	n = size(X, 1)
	m = size(PATH, 1)
	out1 = open("nodeInfo.txt", "w")
    out2 = open("edgeInfo.txt", "w")
	outlabel = open("labelInfo.txt", "w")
	outColor = open("colorInfo.txt", "w")
	outSetting = open("DrawSettings.txt", "w")
	println(outSetting, 5*maxd)
	close(outSetting)
	println(out1, n)
    for i = 1 : n
        println(out1, X[i])
        println(out1, Y[i])
    end
    close(out1)
	for i = 1 : n
		println(outColor, NodeC[labels[i]])
	end
	close(outColor)
	println(out2, m)
	for i = 1 : m
		println(out2, size(PATH[i], 1))
		for j = 1 : size(PATH[i], 1)
			println(out2, PATH[i][j][1])
			println(out2, PATH[i][j][2])
		end
	end
	close(out2)
	for i = 1 : n
		println(outlabel, labels[i])
	end
	close(outlabel)
	outSubline = open("SublineInfo.txt", "w")
	numPath2 = 0
	for i = 1 : size(PATH2, 1)
		if size(PATH2[i], 1) > 0
			numPath2 += 1
		end
	end
	println(outSubline, numPath2)
	for i = 1 : size(PATH2, 1)
		if size(PATH2[i], 1) == 0
			continue
		end
		println(outSubline, size(PATH2[i], 1))
		for j = 1 : size(PATH2[i], 1)
			println(outSubline, PATH2[i][j][1])
			println(outSubline, PATH2[i][j][2])
		end
	end
	close(outSubline)
	run(Cmd(["mpost", "-numbersystem=double", "DrawGraph.mp"]))
	run(Cmd(["mv", "DrawGraph-1.eps", string("../pictures/", na, ".eps")]))
	rm("nodeInfo.txt")
	rm("edgeInfo.txt")
	rm("labelInfo.txt")
	rm("colorInfo.txt")
	rm("DrawSettings.txt")
	rm("SublineInfo.txt")
	rm("DrawGraph.log")
end


text:
using ArchGDAL
using CSV
using DataFrames
using Geodesy
using Images
using ImageTransformations
using Plots

# File-naming-parsing functionality.
get_north(x) = x[1] == 'n' ? parse(Int, x[2:3]) : -parse(Int, x[2:3])
get_east(x) = x[4] == 'e' ? parse(Int, x[5:7]) : -parse(Int, x[5:7])

# Is x inside the bounding box defined by lats / lons?
is_inside(x, lats, lons) = lats[1] <= x.lat <= lats[2] && lons[1] <= x.lon <= lons[2]

# Set these to define the region that the grid covers. Must be set large enough to cover
# all of the elevation that you might be interested in.

# Old GB + Ireland + a-bit-of-France bbox.
# const lons = (-11, 2);
# const lats = (49, 60);

# # Central Europe bbox.
# const lons = (5, 15);
# const lats = (45, 55);

const lons = (-125, -122);
const lats = (47, 49);

# const lon_interval = (-125, -122);
# const lat_interval = (47, 49);

# The NASA data comprises square images of size img_side_length x img_side_length.
# ratio is used to choose lowering the resolution of these images so that the combined grid
# actually fits in memory. A Europe-wide full-resolution image uses about 17GB of memory.
img_side_length = 3601;
ratio = 1 / 8;

# Apply the imresize operation to a fake full-resolution imaget to determine the size of the
# lower-resolution images.
new_size = size(imresize(randn(img_side_length, img_side_length); ratio=ratio));
img_height = new_size[1]
img_width = new_size[2]

# Pre-allocate for the joined-up image data.
elevation_data_mat = fill(
    0.0,
    img_width * (lons[2] - lons[1] + 1),
    img_height * (lats[2] - lats[1] + 1),
);

# Load each of the NASA data files, and plonk them in the correct location in the matrix.
fnames = filter(x -> x[9:11] == "hgt", readdir("scripts/nasa_data_grabbing/nasa_data"));

for fname in fnames
    lat = get_north(fname)
    lon = get_east(fname)
    if is_inside(LLA(lat, lon, 0), lats, lons)

        start_row = img_height * (lat - lats[1]) + 1
        row_range = start_row:(start_row + img_height - 1)

        start_col = img_width * (lon - lons[1]) + 1
        col_range = start_col:(start_col + img_width - 1)

        data = ArchGDAL.read(joinpath("scripts", "nasa_data_grabbing", "nasa_data", fname));

        band = ArchGDAL.getband(data, 1);
        small_band = imresize(band[:, :]; ratio=ratio)

        elevation_data_mat[col_range, row_range] .= reverse(small_band; dims=2)
    end
end;

pixel_lons = range(
    lons[1];
    step=(lons[2] - lons[1] + 1) / size(elevation_data_mat, 1),
    length=size(elevation_data_mat, 1),
);
pixel_lats = range(
    lats[1];
    step=(lats[2] - lats[1] + 1) / size(elevation_data_mat, 2),
    length=size(elevation_data_mat, 2),
);

# Functionality to determine the value of the closest pixel to the requested lat-lon.
function nearest_point(elevation_data, lat::Real, lon::Real)
    return elevation_data[
        nearest_point_idx(pixel_lons, lon),
        nearest_point_idx(pixel_lats, lat),
    ]
end

function nearest_point_idx(xs::StepRangeLen, x::Real)
    return convert(Int, ceil((x - first(xs) + 1e-9) / (last(xs) - first(xs)) * length(xs)))
end

# # Plot the matrix as an image.
# min_val = minimum(elevation_data_mat);
# max_val = maximum(elevation_data_mat);
# img = Gray.((elevation_data_mat .- min_val) ./ (max_val - min_val));
# Plots.plot(reverse(img; dims=2)')

# log_img = log.(Float64.(img) .+ 1);
# min_log_val = minimum(log_img);
# max_log_val = maximum(log_img);
# new_img = Gray.((log_img .- min_log_val) ./ (max_log_val - min_log_val));
# Plots.plot(reverse(new_img; dims=2)')

# heatmap(log.(elevation_data_mat .- minimum(elevation_data_mat) .+ 1))