text:
export realizer, dimension, realize_poset

using JuMP

_distinct(a,b,c) = (a!=b) && (a!=c) && (b!=c)

"""
`realizer(P::SimplePoset,d::Int)` creates a realizer of `P` using `d` linear
extensions or throws an error if none exists. The output is an `n`-by-`d` matrix
whose columns give the linear extensions. The first element in each column is the
bottom element of that linear extension.
"""
function realizer(P::SimplePoset{T}, d::Int) where T
    # MOD = Model(solver=SimpleGraphAlgorithms._SOLVER())
    MOD = Model(with_optimizer(SimpleGraphAlgorithms._SOLVER.Optimizer;
        SimpleGraphAlgorithms._OPTS...))

    VV = elements(P)
    n = length(VV)
    # x[u,v,t] == 1 means u<v in L_t

    @variable(MOD, x[u=VV, v=VV, t=1:d] , Bin)

    # for all i, x[i,i,t] is zero
    for u in VV
        for t in 1:d
            @constraint(MOD,x[u,u,t] == 0)
        end
    end

    # exactly one of x[u,v,t] or x[v,u,t] is 1
    for i=1:n-1
        for j=i+1:n
            u = VV[i]
            v = VV[j]
            for t=1:d
                @constraint(MOD,x[u,v,t]+x[v,u,t] == 1)
            end
        end
    end

    # if u<v x[u,v,t] == 1
    for u in VV
        for v in VV
            if has(P,u,v)
                for t in 1:d
                    @constraint(MOD, x[u,v,t] == 1)
                end
            end
        end
    end

    # if u and v are incomparable, sum(X[u,v,t]) > 0
    for i=1:n-1
        for j=i+1:n
            u = VV[i]
            v = VV[j]
            if !has(P,u,v) && !has(P,v,u)
                @constraint(MOD, sum(x[u,v,t] for t in 1:d) >= 1)
                @constraint(MOD, sum(x[v,u,t] for t in 1:d) >= 1)
            end
        end
    end

    # ensure L_t is transitive (so linear)
    for t=1:d
        for u in VV
            for v in VV
                for w in VV
                    if _distinct(u,v,w)
                        @constraint(MOD, x[u,w,t] >= x[u,v,t]+x[v,w,t]-1)
                    end
                end
            end
        end
    end

    optimize!(MOD)
    status = Int(termination_status(MOD))

    if status != 1
        error("This poset has dimension greater than $d; no realizer found.")
    end

    X = value.(x)

    PP = [ SimplePoset{T}() for t =1:d]
    for t=1:d
        for u in VV
            for v in VV
                if X[u,v,t] > 0.5
                    add!(PP[t],u,v)
                end
            end
        end
    end

    LL = [ linear_extension(PP[t]) for t=1:d ]

    result = Array{T,2}(undef,n,d)
    for t=1:d
        for i=1:n
            result[i,t] = LL[t][i]
        end
    end

    return result
end

"""
`realize_poset(R::Array{T,2})` creates a poset from a realizer.
The columns of `R` are the linear extensions of some poset; this
function returns that poset.
"""
function realize_poset(R::Array{T,2})::SimplePoset{T} where T
    n,d = size(R)
    P = make_linear_order(R[:,1])
    for j=2:d
        L = make_linear_order(R[:,j])
        P = intersect(L,P)
    end
    return P
end

function make_linear_order(lst::Array{T,1})::SimplePoset{T} where T
    P = SimplePoset{T}()
    n = length(lst)
    for i=1:n-1
        add!(P,lst[i],lst[i+1])
    end
    return P
end

"""
`dimension(P::SimplePoset, verbose=false)` returns the order-theoretic
dimension of the poset `P`. Set `verbose` to `true` to see more information
as the work is done.
"""
function dimension(P::SimplePoset, verb::Bool = false)::Int
    n = card(P)
    if n==0
        return 0
    end

    if length(incomparables(P))==0  # it's a chain
        return 1
    end

    lb = 2

    ub1 = Int(floor(n/2))
    ub2 = width(P)
    ub = max(ub1,ub2)

    return dimension_work(P, lb, ub, verb)
end

function dimension_work(P::SimplePoset, lb::Int, ub::Int, verb::Bool)::Int
    if verb
        print("$lb <= dim(P) <= $ub\t")
    end

    if lb == ub
        if verb
            println("and we're done")
        end
        return lb
    end

    mid = Int(floor((ub+lb)/2))

    if verb
        print("looking for a $mid realizer\t")
    end

    try
        R = realizer(P,mid)
        if verb
            println("confirmed")
        end
        return dimension_work(P,lb,mid,verb)
    catch
    end
    if verb
        println("none exists")
    end
    return dimension_work(P,mid+1,ub,verb)
end


text:
"""
    tv(
        f::AbstractArray{<:AbstractFloat, N ∈ (3, 4)},
        mask::AbstractArray{Bool, 3},
        vsz::NTuple{3, Real};
        W::Union{Nothing, AbstractArray{<:AbstractFloat, M ∈ (3, N)}} = nothing,
        Wtv::Union{Nothing, AbstractArray{<:AbstractFloat, M ∈ (3, N)}} = nothing,
        pad::NTuple{3, Integer} = (0, 0, 0),
        Dkernel::Symbol = :k,
        bdir::NTuple{3, Real} = (0, 0, 1),
        lambda::Real = 1e-3,
        rho::Real = 100*lambda,
        mu::Real = 1,
        tol::Real = 1e-3,
        maxit::Integer = 250,
        verbose::Bool = false,
    ) -> typeof(similar(f))

Total variation deconvolution using ADMM [1].

### Arguments
- `f::AbstractArray{<:AbstractFloat, N ∈ (3, 4)}`: unwrapped (multi-echo) local
    field/phase
- `mask::AbstractArray{Bool, 3}`: binary mask of region of interest
- `vsz::NTuple{3, Real}`: voxel size

### Keywords
- `W::Union{Nothing, AbstractArray{<:AbstractFloat, M ∈ (3, N)}} = nothing`:
    data fidelity weights
- `Wtv::Union{Nothing, AbstractArray{<:AbstractFloat, M ∈ (3, 5)}} = nothing`:
    total variation weights
    - `M = 3`: same weights for all three gradient directions and all echoes
    - `M = 4 = N`: same weights for all three gradient directions, different weights for echoes
    - `M = 5, (size(Wtv)[4,5] = [1 or N, 3]`: different weights for each gradient direction
- `pad::NTuple{3, Integer} = (0, 0, 0)`: zero padding array
    - `< 0`: no padding
    - `≥ 0`: minimum padding to fast fft size
- `bdir::NTuple{3, Real} = (0, 0, 1)`: unit vector of B field direction
- `Dkernel::Symbol = :k`: dipole kernel method
- `lambda::Real = 1e-3`: regularization parameter
- `rho::Real = 100*lambda`: Lagrange multiplier penalty parameter
- `mu::Real = 1`: Lagrange multiplier penalty parameter (unused if `W = nothing`)
- `tol::Real = 1e-3`: stopping tolerance
- `maxit::Integer = 250`: maximum number of iterations
- `verbose::Bool = false`: print convergence information

### Returns
- `typeof(similar(f))`: susceptibility map

### References
[1] Bilgic B, Fan AP, Polimeni JR, Cauley SF, Bianciardi M, Adalsteinsson E,
    Wald LL, Setsompop K. Fast quantitative susceptibility mapping with
    L1‐regularization and automatic parameter selection.
    Magnetic resonance in medicine. 2014 Nov;72(5):1444-59.
"""
function tv(
    f::AbstractArray{T, N},
    mask::AbstractArray{Bool, 3},
    vsz::NTuple{3, Real};
    W::Union{Nothing, AbstractArray{<:AbstractFloat}} = nothing,
    Wtv::Union{Nothing, AbstractArray{<:AbstractFloat}} = nothing,
    pad::NTuple{3, Integer} = (0, 0, 0),
    Dkernel::Symbol = :k,
    bdir::NTuple{3, Real} = (0, 0, 1),
    lambda::Real = 1e-3,
    rho::Real = 100*lambda,
    mu::Real = 1,
    tol::Real = 1e-3,
    maxit::Integer = 250,
    verbose::Bool = false,
) where {T, N}
    N ∈ (3, 4) || throw(ArgumentError("arrays must be 3d or 4d, got $(N)d"))
    return _tv!(
        tzero(f), f, mask, vsz, W, Wtv, pad, Dkernel, bdir,lambda,
        rho, mu, tol, maxit, verbose
    )
end

function _tv!(
    x::AbstractArray{<:AbstractFloat, N},
    f::AbstractArray{T, N},
    mask::AbstractArray{Bool, 3},
    vsz::NTuple{3, Real},
    W::Union{Nothing, AbstractArray{<:AbstractFloat}},
    Wtv::Union{Nothing, AbstractArray{<:AbstractFloat}},
    pad::NTuple{3, Integer},
    Dkernel::Symbol,
    bdir::NTuple{3, Real},
    lambda::Real,
    rho::Real,
    mu::Real,
    tol::Real,
    maxit::Integer,
    verbose::Bool,
) where {T, N}
    N ∈ (3, 4) || throw(ArgumentError("arrays must be 3d or 4d, got $(N)d"))

    size(x) == size(f) || throw(DimensionMismatch())
    size(mask) == size(f)[1:3] || throw(DimensionMismatch())

    if W !== nothing
        ndims(W) ∈ (3, 4) || throw(ArgumentError("arrays must be 3d or 4d, got $(M)d"))
        size(W)[1:3] == size(f)[1:3] || throw(DimensionMismatch())
        ndims(W) == 3 || size(W, 4) == size(f, 4) || throw(DimensionMismatch())
    end

    if Wtv !== nothing
        ndims(Wtv) ∈ (3, 4, 5) || throw(ArgumentError("array Wtv must be 3d, 4d, or 5d, got $(M)d"))
        size(Wtv)[1:3] == size(f)[1:3] || throw(DimensionMismatch())
        ndims(Wtv) == 3 || size(Wtv, 4) ∈ (1, size(f, 4)) || throw(DimensionMismatch())
        ndims(Wtv) ∈ (3, 4) || size(Wtv, 5) == 3 || throw(DimensionMismatch())
    end

    Dkernel ∈ (:k, :kspace, :i, :ispace) ||
        throw(ArgumentError("Dkernel must be one of :k, :kspace, :i, :ispace, got :$(Dkernel)"))

    # convert scalars
    _zero = zero(T)

    ρ = convert(T, rho)
    μ = convert(T, mu)
    λ = convert(T, lambda)
    ϵ = convert(T, tol)

    iρ = inv(ρ)
    λiρ = λ*iρ

    # pad to fast fft size
    xp = padfastfft(@view(f[:,:,:,1]), pad, rfft=true)
    m = padfastfft(mask, pad, rfft=true)

    # initialize variables and fft
    sz0 = size(mask)
    sz  = size(m)
    sz_ = (sz[1]>>1 + 1, sz[2], sz[3])

    x0 = similar(xp)

    D  = Array{T}(undef, sz_)           # dipole kernel
    L  = Array{T}(undef, sz_)           # laplace kernel
    iA = Array{T}(undef, sz_)           # 1 / (D'D - rho*Δ)

    X̂ = Array{complex(T)}(undef, sz_)   # in-place rfft var
    F̂ = Array{complex(T)}(undef, sz_)   # pre-computed rhs

    FFTW.set_num_threads(num_cores())
    P = plan_rfft(xp)
    iP = inv(P)

    # admm
    ux = similar(xp)
    uy = similar(xp)
    uz = similar(xp)

    dx = similar(xp)
    dy = similar(xp)
    dz = similar(xp)

    if W !== nothing
        iAw = similar(xp)
        fw = similar(xp)
        uw = similar(xp)
    end

    if Wtv !== nothing
        if ndims(Wtv) == 3
            λiρWx = padarray!(similar(xp), Wtv)
            λiρWy = λiρWx
            λiρWz = λiρWx
        elseif ndims(Wtv) == 4
            λiρWx = similar(xp)
            λiρWy = λiρWx
            λiρWz = λiρWx
        elseif size(Wtv, 4) == 1
            λiρWx = padarray!(similar(xp), @view(Wtv[:,:,:,1,1]))
            λiρWx = padarray!(similar(xp), @view(Wtv[:,:,:,1,2]))
            λiρWx = padarray!(similar(xp), @view(Wtv[:,:,:,1,3]))
        else
            λiρWx = similar(xp)
            λiρWy = similar(xp)
            λiρWz = similar(xp)
        end
    end

    # get kernels
    D = _dipole_kernel!(D, F̂, xp, sz0, vsz, bdir, P, Dkernel, :rfft)
    L = _laplace_kernel!(L, F̂, xp, vsz, P, negative=true)

    @inbounds for t in axes(f, 4)
        if verbose && size(f, 4) > 1
            @printf("Echo: %d/%d\n", t, size(f, 4))
        end

        xp = padarray!(xp, @view(f[:, :, :, t]))

        # set up constant rhs and inverse lhs
        if W !== nothing
            # iA  = (μ*D^H*D - ρ*Δ)^-1
            # iAw = μ / (W^t*W + μ)
            # fw  = iAw * W^t * f
            F̂ = tfill!(F̂, 0)

            @batch for I in eachindex(iA)
                a = μ*conj(D[I])*D[I] + ρ*L[I]
                iA[I] = ifelse(iszero(a), _zero, inv(a))
            end

            fw = padarray!(fw, @view(W[:,:,:,min(t, end)]))

            @batch for I in eachindex(fw)
                w = fw[I]
                ia = inv(muladd(w, w, μ))
                iAw[I] = μ*ia
                fw[I] = uw[I] = ia*w*xp[I]
            end

        else
            # iA = (D^H*D - ρ*Δ)^-1
            # F̂  = iA * D^H*f
            X̂ = mul!(X̂, P, xp)

            @batch for I in eachindex(F̂)
                a = conj(D[I])*D[I] + ρ*L[I]
                if iszero(a)
                    F̂[I] = _zero
                    iA[I] = _zero
                else
                    ia = inv(a)
                    F̂[I] = ia * conj(D[I]) * X̂[I]
                    iA[I] = ia
                end
            end
        end

        # total variation weights
        if Wtv !== nothing
            if ndims(Wtv) == 4
                λiρWx = padarray!(λiρWx, @view(Wtv[:,:,:,t]))
            elseif ndims(Wtv) == 5 && size(Wtv, 4) > 1
                λiρWx = padarray!(λiρWx, @view(Wtv[:,:,:,t,1]))
                λiρWy = padarray!(λiρWy, @view(Wtv[:,:,:,t,2]))
                λiρWz = padarray!(λiρWz, @view(Wtv[:,:,:,t,3]))
            end
        end

        ux = tfill!(ux, _zero)
        uy = tfill!(uy, _zero)
        uz = tfill!(uz, _zero)
        dx, dy, dz = _gradfp!(dx, dy, dz, xp, vsz)

        if verbose
            @printf("\n iter\t  ||x-xprev||/||x||\n")
        end

        for i in 1:maxit
            x0, xp = xp, x0

            ##################################################################
            # x - subproblem
            ##################################################################
            # d = z - u
            xp = _gradfp_adj!(xp, dx, dy, dz, vsz)

            X̂ = mul!(X̂, P, xp)

            @batch for I in eachindex(X̂)
                X̂[I] *= ρ*iA[I]
                X̂[I] += F̂[I]
            end

            if W !== nothing
                F̂ = _tcopyto!(F̂, X̂) # real ifft overwrites input
            end

            xp = mul!(xp, iP, X̂)

            ##################################################################
            # convergence check
            ##################################################################
            # [1] ndx = norm(x - xprev)
            # [2] nx  = norm(x)
            @batch threadlocal=zeros(T, 2)::Vector{T} for I in eachindex(xp)
                a, b = xp[I], x0[I]
                threadlocal[1] = muladd(a-b, a-b, threadlocal[1])
                threadlocal[2] = muladd(a, a, threadlocal[2])
            end
            ndx, nx = sqrt.(sum(threadlocal::Vector{Vector{T}}))

            if verbose
                @printf("%3d/%d\t    %.4e\n", i, maxit, ndx/nx)
            end

            if ndx < ϵ*nx || i == maxit
                break
            end

            ##################################################################
            # z - subproblem and Lagrange multiplier update
            ##################################################################
            dx, dy, dz = _gradfp!(dx, dy, dz, xp, vsz)

            # z = shrink(∇x + u, λ/ρ)   z-subproblem
            # u = u + ∇x - z            Lagrange multiplier update
            # d = z - u                 pre-compute for x-problem
            if Wtv !== nothing
                @batch for I in eachindex(ux)
                    λiρ = λiρWx[I]
                    u = ux[I] + dx[I]
                    z = ifelse(abs(u) ≤ λiρ, _zero, copysign(abs(u)-λiρ, u))
                    ux[I] = u - z
                    dx[I] = z - u + z
                end

                @batch for I in eachindex(uy)
                    λiρ = λiρWy[I]
                    u = uy[I] + dy[I]
                    z = ifelse(abs(u) ≤ λiρ, _zero, copysign(abs(u)-λiρ, u))
                    uy[I] = u - z
                    dy[I] = z - u + z
                end

                @batch for I in eachindex(uz)
                    λiρ = λiρWz[I]
                    u = uz[I] + dz[I]
                    z = ifelse(abs(u) ≤ λiρ, _zero, copysign(abs(u)-λiρ, u))
                    uz[I] = u - z
                    dz[I] = z - u + z
                end

            else
                @batch for I in eachindex(ux)
                    u = ux[I] + dx[I]
                    z = ifelse(abs(u) ≤ λiρ, _zero, copysign(abs(u)-λiρ, u))
                    ux[I] = u - z
                    dx[I] = z - u + z
                end

                @batch for I in eachindex(uy)
                    u = uy[I] + dy[I]
                    z = ifelse(abs(u) ≤ λiρ, _zero, copysign(abs(u)-λiρ, u))
                    uy[I] = u - z
                    dy[I] = z - u + z
                end

                @batch for I in eachindex(uz)
                    u = uz[I] + dz[I]
                    z = ifelse(abs(u) ≤ λiρ, _zero, copysign(abs(u)-λiρ, u))
                    uz[I] = u - z
                    dz[I] = z - u + z
                end
            end

            ##################################################################
            # F̂ - update
            ##################################################################
            if W !== nothing
                # F̂ = μ*D^H*(z - u) / (μ*D^H*D - ρ*Δ)
                # z = (W*f + μ*(D*x + u)) / (W^t*W + μ)
                # u = u + D*x - z
                @batch for I in eachindex(F̂)
                    F̂[I] *= D[I]
                end

                x0 = mul!(x0, iP, F̂)
                @batch for I in eachindex(x0)
                    u = x0[I] + uw[I]
                    z = fw[I] + iAw[I]*u
                    uw[I] = u - z
                    x0[I] = z - u + z
                end

                F̂ = mul!(F̂, P, x0)
                @batch for I in eachindex(F̂)
                    F̂[I] *= μ*conj(D[I])*iA[I]
                end
            end
        end

        unpadarray!(@view(x[:,:,:,t]), xp)

        if verbose
            println()
        end

    end

    return x
end

ext: jl


text:
filename() = "dayXX_sample.txt"

function read_data()

end

function part1()
end

function part2()
end


text:
include("splinesSerial.jl")

controlpoints=[
   [[ 0,0,0],[0 ,3  ,4],[0,6,3],[0,10,0]],
   [[ 3,0,2],[2 ,2.5,5],[3,6,5],[4,8,2]],
   [[ 6,0,2],[8 ,3 , 5],[7,6,4.5],[6,10,2.5]],
   [[10,0,0],[11,3  ,4],[11,6,3],[10,9,0]]]

bicubic = larBicubicSurface(controlpoints)
point = [0.7,0.2]
result = bicubic(point)


text:
using Printf

prog = Base.basename(Base.source_path())

println(prog, "'s command-line arguments are:")
for s in ARGS
    println("    ", s)
end

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
function construct_device!(ps_m::CanonicalModel,
                           device::Type{B},
                           device_formulation::Type{Br},
                           system_formulation::Type{StandardPTDFForm},
                           sys::PSY.System,
                           time_range::UnitRange{Int64};
                           kwargs...) where {Br <: AbstractLineForm,
                                             B <: PSY.Branch}

    line_rate_constraints(ps_m, sys.branches, device_formulation, system_formulation, time_range)

    return

end

function construct_device!(ps_m::CanonicalModel,
                            device::Type{B},
                            device_formulation::Type{Br},
                            system_formulation::Type{S},
                            sys::PSY.System,
                            time_range::UnitRange{Int64};
                            kwargs...) where {Br <: AbstractLineForm,
                                              B <: PSY.Branch,
                                              S <: PM.AbstractPowerFormulation}

    # This code is meant to do nothing and will have a constructor once the branch formulations are developed

    return

end


text:
using CSV, Distributions, HypothesisTests

data = CSV.read("../data/machine1.csv", header=false)[:,1]
xBar, n = mean(data), length(data)
s = std(data)
mu0 = 52.2

testStatistic = ( xBar - mu0 ) / ( s/sqrt(n) )
pVal = 2*ccdf(TDist(n-1), abs(testStatistic))

println("Manually calculated test statistic: ", testStatistic)
println("Manually calculated p-value: ", pVal,"\n")
OneSampleTTest(data, mu0)

max_stars_repo_name: Yoshinobu-Ishizaki/StatsWithJuliaBook

max_issues_repo_name: Yoshinobu-Ishizaki/StatsWithJuliaBook

max_forks_repo_name: Yoshinobu-Ishizaki/StatsWithJuliaBook


text:
# ------------------------------------------------------------------------------------------------ #
#                                        Windowing functions                                       #
# ------------------------------------------------------------------------------------------------ #

function slidingWindow(x::AbstractVector, width::Int=length(x)÷20, overlap::Int=0, windowFunc::Function=rect)
    # windowFunc must be vectorised
    # Construct sub-timeseries
    X = hcat([x[Array(1:width).+a] for a in 0:(width-overlap):(length(x)-width)]...)
    idxs = 1:size(X, 1):(size(X, 2)*size(X, 1)+1)
    # Apply the window function to each sub-series
    (mapslices(windowFunc, X, dims=1), idxs)
end
function slidingWindow(width::Int, args...)
    f(x) = slidingWindow(x, width, args...)
end
export slidingWindow

function slidingwindow(x::AbstractVector, width::Int=length(x)÷20, overlap::Int=0, windowFunc::Function=rect)
    # windowFunc must be vectorised
    # Construct sub-timeseries
    X = hcat([x[Array(1:width).+a] for a in 0:(width-overlap):(length(x)-width)]...)
    idxs = 1:size(X, 1):(size(X, 2)*size(X, 1)+1)
    # Apply the window function to each sub-series
    mapslices(windowFunc, X, dims=1)
end
function slidingwindow(width::Int, args...)
    f(x) = slidingwindow(x, width, args...)
end
export slidingwindow

# ------------------------------------------------------------------------------------------------ #
#                                         Window functions                                         #
# ------------------------------------------------------------------------------------------------ #
rect(x::AbstractVector) = x
export rect


text:
# dummy type implementing the Data trait
struct DummyData{𝒟,𝒱} <: Data
  domain::𝒟
  values::𝒱
end

TopOptDash.domain(data::DummyData) = getfield(data, :domain)

function TopOptDash.values(data::DummyData, rank=nothing)
  domain = getfield(data, :domain)
  values = getfield(data, :values)
  r = isnothing(rank) ? paramdim(domain) : rank
  haskey(values, r) ? values[r] : nothing
end

TopOptDash.constructor(::Type{D}) where {D<:DummyData} = DummyData


text:
module EnsembleKalmanProcessModule

using ..ParameterDistributionStorage
using ..DataStorage

using Random
using Statistics
using Distributions
using LinearAlgebra
using DocStringExtensions

export EnsembleKalmanProcess, Inversion, Sampler, Unscented
export MiniBatchKalmanProcess
export get_u, get_g, get_obs, get_obs_cov
export get_u_prior, get_u_final, get_u_mean_final, get_g_final
export get_N_iterations, get_error, get_obs_final, get_obs_cov_final
export construct_initial_ensemble
export compute_error!
export update_ensemble!
export find_ekp_stepsize
export Gaussian_2d

abstract type Process end

"""
    Inversion <: Process

An ensemble Kalman Inversion process
"""
struct Inversion <: Process end

"""
    Sampler{FT<:AbstractFloat,IT<:Int} <: Process

An ensemble Kalman Sampler process
"""
struct Sampler{FT<:AbstractFloat} <: Process
  ""
  prior_mean::Vector{FT}
  ""
  prior_cov::Array{FT, 2}
end

"""
    Unscented{FT<:AbstractFloat, IT<:Int} <: Process

An unscented Kalman Inversion process
"""
mutable struct  Unscented{FT<:AbstractFloat, IT<:Int} <: Process
    "a vector of arrays of size N_parameters containing the mean of the parameters (in each uki iteration a new array of mean is added)"
    u_mean::Vector{Array{FT, 1}}
    "a vector of arrays of size (N_parameters x N_parameters) containing the covariance of the parameters (in each uki iteration a new array of cov is added)"
    uu_cov::Vector{Array{FT, 2}}
    "a vector of arrays of size N_y containing the predicted observation (in each uki iteration a new array of predicted observation is added)"
    obs_pred::Vector{Array{FT, 1}}
    "weights in UKI"
    c_weights::Array{FT, 1}
    mean_weights::Array{FT, 1}
    cov_weights::Array{FT, 1}
    "covariance of the artificial evolution error"
    Σ_ω::Array{FT, 2}
    "covariance of the artificial observation error"
    Σ_ν_scale::FT
    "regularization parameter"
    α_reg::FT
    "regularization vector"
    r::Array{FT, 1}
    "update frequency"
    update_freq::IT
    "current iteration number"
    iter::IT
end

abstract type KalmanProcessObject end

"""
    EnsembleKalmanProcess{FT<:AbstractFloat, IT<:Int}

Structure that is used in Ensemble Kalman processes

#Fields
$(DocStringExtensions.FIELDS)
"""
struct EnsembleKalmanProcess{FT<:AbstractFloat, IT<:Int, P<:Process} <: KalmanProcessObject
    "Array of stores for parameters (u), each of size [N_par × N_ens]"
    u::Array{DataContainer{FT}}
    "vector of the observed vector size [N_obs]"
    obs_mean::Vector{FT}
    "covariance matrix of the observational noise, of size [N_obs × N_obs]"
    obs_noise_cov::Array{FT, 2}
    "ensemble size"
    N_ens::IT
    "Array of stores for forward model outputs, each of size  [N_obs × N_ens]"
    g::Array{DataContainer{FT}}
    "vector of errors"
    err::Vector{FT}
    "vector of timesteps used in each EK iteration"
    Δt::Vector{FT}
    "the particular EK process (`Inversion` or `Sampler`)"
    process::P
end

# outer constructors
function EnsembleKalmanProcess(params::Array{FT, 2},
                               obs_mean,
                               obs_noise_cov::Array{FT, 2},
                               process::P;
                               Δt=FT(1)) where {FT<:AbstractFloat, P<:Process}

    #initial parameters stored as columns
    init_params=DataContainer(params, data_are_columns=true)
    # ensemble size
    N_ens = size(init_params,2) #stored with data as columns
    IT = typeof(N_ens)
    #store for model evaluations
    g=[]
    # error store
    err = FT[]
    # timestep store
    Δt = Array([Δt])

    EnsembleKalmanProcess{FT, IT, P}([init_params], obs_mean, obs_noise_cov, N_ens, g,
                                     err, Δt, process)
end

struct MiniBatchKalmanProcess{FT<:AbstractFloat, IT<:Int, P<:Process} <: KalmanProcessObject
    "Array of stores for parameters (u), each of size [N_par × N_ens]"
    u::Array{DataContainer{FT}}
    "Array of stores for observations, each of size [N_obs_i]"
    obs_mean::Vector{Array{FT, 1}}
    "Array of stores for covariance matrix of the observational noise, each of size [N_obs_i × N_obs_i]"
    obs_noise_cov::Array{DataContainer{FT}}
    "ensemble size"
    N_ens::IT
    "Array of stores for forward model outputs, each of size  [N_obs × N_ens]"
    g::Array{DataContainer{FT}}
    "vector of errors"
    err::Vector{FT}
    "vector of timesteps used in each MBK iteration"
    Δt::Vector{FT}
    "the particular MBK process (`Inversion` or `Sampler`)"
    process::P
end

# outer constructors
function MiniBatchKalmanProcess(params::Array{FT, 2},
                               process::P;
                               Δt=FT(1)) where {FT<:AbstractFloat, P<:Process}

    #initial parameters stored as columns
    init_params=DataContainer(params, data_are_columns=true)
    # ensemble size
    N_ens = size(init_params,2) #stored with data as columns
    IT = typeof(N_ens)
    #store for model evaluations
    g=[]
    # store for observations
    obs_mean = []
    obs_noise_cov = []
    # error store
    err = FT[]
    # timestep store
    Δt = Array([Δt])

    MiniBatchKalmanProcess{FT, IT, P}([init_params], obs_mean, obs_noise_cov, N_ens, g,
                                     err, Δt, process)
end

"""
    get_u(ekp::EnsembleKalmanProcess, iteration::IT; return_array=true) where {IT <: Integer}

Get for the EKI iteration. Returns a DataContainer object unless array is specified.
"""
function get_u(ekp::KalmanProcessObject, iteration::IT; return_array=true) where {IT <: Integer}
    return  return_array ? get_data(ekp.u[iteration]) : ekp.u[iteration]
end

"""
    get_g(ekp::EnsembleKalmanProcess, iteration::IT; return_array=true) where {IT <: Integer}

Get for the EKI iteration. Returns a DataContainer object unless array is specified.
"""
function get_g(ekp::KalmanProcessObject, iteration::IT; return_array=true) where {IT <: Integer}
    return return_array ? get_data(ekp.g[iteration]) : ekp.g[iteration]
end

"""
    get_u(ekp::EnsembleKalmanProcess; return_array=true)

Get for the EKI iteration. Returns a DataContainer object unless array is specified.
"""
function get_u(ekp::KalmanProcessObject; return_array=true) where {IT <: Integer}
    N_stored_u = get_N_iterations(ekp)+1
    return [get_u(ekp, it, return_array=return_array) for it in 1:N_stored_u]
end

"""
    get_g(ekp::EnsembleKalmanProcess; return_array=true)

Get for the EKI iteration. Returns a DataContainer object unless array is specified.
"""
function get_g(ekp::KalmanProcessObject; return_array=true) where {IT <: Integer}
    N_stored_g = get_N_iterations(ekp)
    return [get_g(ekp, it, return_array=return_array) for it in 1:N_stored_g]
end

"""
    get_u_final(ekp::EnsembleKalmanProcess, return_array=true)

Get the final or prior iteration of parameters or model ouputs, returns a DataContainer Object if return_array is false.
"""
function get_u_final(ekp::KalmanProcessObject; return_array=true)
    return return_array ? get_u(ekp,size(ekp.u)[1]) : ekp.u[end]
end

"""
    get_u_prior(ekp::EnsembleKalmanProcess, return_array=true)

Get the final or prior iteration of parameters or model ouputs, returns a DataContainer Object if return_array is false.
"""

function get_u_prior(ekp::KalmanProcessObject; return_array=true)
    return return_array ? get_u(ekp,1) : ekp.u[1]
end

"""
    get_g_final(ekp::EnsembleKalmanProcess, return_array=true)

Get the final or prior iteration of parameters or model ouputs, returns a DataContainer Object if return_array is false.
"""

function get_g_final(ekp::KalmanProcessObject; return_array=true)
    return return_array ? get_g(ekp,size(ekp.g)[1]) : ekp.g[end]
end

function get_obs(ekp::MiniBatchKalmanProcess, iteration::IT) where {IT <: Integer}
    return ekp.obs_mean[iteration]
end

function get_obs(ekp::MiniBatchKalmanProcess; return_array=true) where {IT <: Integer}
    N_stored_obs = get_N_iterations(ekp)
    return [get_obs(ekp, it) for it in 1:N_stored_obs]
end

function get_obs_final(ekp::MiniBatchKalmanProcess)
    return ekp.obs_mean[end]
end

function get_obs_cov(ekp::MiniBatchKalmanProcess, iteration::IT; return_array=true) where {IT <: Integer}
    return return_array ? get_data(ekp.obs_noise_cov[iteration]) : ekp.obs_noise_cov[iteration]
end

function get_obs_cov(ekp::MiniBatchKalmanProcess; return_array=true) where {IT <: Integer}
    N_stored_obs = get_N_iterations(ekp)
    return [get_obs_cov(ekp, it, return_array=return_array) for it in 1:N_stored_obs]
end

function get_obs_cov_final(ekp::MiniBatchKalmanProcess; return_array=true)
    return return_array ? get_obs_cov(ekp,size(ekp.obs_noise_cov)[1]) : ekp.obs_noise_cov[end]
end

"""
    get_N_iterations(ekp::EnsembleKalmanProcess

get number of times update has been called (equals size(g), or size(u)-1)
"""
function get_N_iterations(ekp::KalmanProcessObject)
    return size(ekp.u)[1] - 1
end

"""
    construct_initial_ensemble(prior::ParameterDistribution, N_ens::IT; rng_seed=42) where {IT<:Int}

Construct the initial parameters, by sampling N_ens samples from specified
prior distribution. Returned with parameters as columns
"""
function construct_initial_ensemble(prior::ParameterDistribution, N_ens::IT; rng_seed=42) where {IT<:Int}
    # Ensuring reproducibility of the sampled parameter values
    Random.seed!(rng_seed)
    parameters = sample_distribution(prior, N_ens) #of size [dim(param space) N_ens]
    return parameters
end

function compute_error!(ekp::EnsembleKalmanProcess)
    mean_g = dropdims(mean(get_g_final(ekp), dims=2), dims=2)
    diff = ekp.obs_mean - mean_g
    X = ekp.obs_noise_cov \ diff # diff: column vector
    newerr = dot(diff, X)
    push!(ekp.err, newerr)
end

function compute_error!(ekp::MiniBatchKalmanProcess)
    mean_g = dropdims(mean(get_g_final(ekp), dims=2), dims=2)
    obs_mean_ = get_obs_final(ekp)
    obs_cov_ = get_obs_cov_final(ekp)
    diff = obs_mean_ - mean_g
    X = obs_cov_ \ diff # diff: column vector
    newerr = dot(diff, X)
    push!(ekp.err, newerr)
end

function get_error(ekp::KalmanProcessObject)
    return ekp.err
end

"""
   find_ekp_stepsize(ekp::EnsembleKalmanProcess{FT, IT, Inversion}, g::Array{FT, 2}; cov_threshold::FT=0.01) where {FT}

Find largest stepsize for the EK solver that leads to a reduction of the determinant of the sample
covariance matrix no greater than cov_threshold.
"""
function find_ekp_stepsize(ekp::EnsembleKalmanProcess{FT, IT, Inversion}, g::Array{FT,2}; cov_threshold::FT=0.01) where {FT, IT}
    accept_stepsize = false
    if !isempty(ekp.Δt)
        Δt = deepcopy(ekp.Δt[end])
    else
        Δt = FT(1)
    end
    # final_params [N_par × N_ens]
    cov_init = cov(get_u_final(ekp), dims=2)
    while accept_stepsize == false
        ekp_copy = deepcopy(ekp)
        update_ensemble!(ekp_copy, g, Δt_new=Δt)
        cov_new = cov(get_u_final(ekp_copy), dims=2)
        if det(cov_new) > cov_threshold * det(cov_init)
            accept_stepsize = true
        else
            Δt = Δt/2
        end
    end

    return Δt

end

"""
    update_ensemble!(ekp::EnsembleKalmanProcess{FT, IT, <:Process}, g::Array{FT,2} cov_threshold::FT=0.01, Δt_new=nothing) where {FT, IT}

Updates the ensemble according to which type of Process we have. Model outputs g need to be a N_obs × N_ens array (i.e data are columms)
"""
function update_ensemble!(ekp::EnsembleKalmanProcess{FT, IT, Inversion},
                          g::Array{FT,2};
                          cov_threshold::FT=0.01,
                          Δt_new=nothing,
                          deterministic_forward_map=true,) where {FT, IT}

    # Update follows eqns. (4) and (5) of Schillings and Stuart (2017)

    #catch works when g non-square
    if !(size(g)[2] == ekp.N_ens)
         throw(DimensionMismatch("ensemble size in EnsembleKalmanProcess and g do not match, try transposing g or check ensemble size"))
    end

    # u: N_par × N_ens
    # g: N_obs × N_ens
    u = get_u_final(ekp)
    N_obs = size(g, 1)

    cov_init = cov(u, dims=2)
    cov_ug = cov(u,g, dims = 2, corrected=false) # [N_par × N_obs]
    cov_gg = cov(g,g, dims = 2, corrected=false) # [N_obs × N_obs]

    if !isnothing(Δt_new)
        push!(ekp.Δt, Δt_new)
    elseif isnothing(Δt_new) && isempty(ekp.Δt)
        push!(ekp.Δt, FT(1))
    else
        push!(ekp.Δt, ekp.Δt[end])
    end

    # Scale noise using Δt
    scaled_obs_noise_cov = ekp.obs_noise_cov / ekp.Δt[end]
    noise = rand(MvNormal(zeros(N_obs), scaled_obs_noise_cov), ekp.N_ens)

    # Add obs_mean (N_obs) to each column of noise (N_obs × N_ens) if
    # G is deterministic
    y = deterministic_forward_map ? (ekp.obs_mean .+ noise) : (ekp.obs_mean .+ zero(noise))

    # N_obs × N_obs \ [N_obs × N_ens]
    # --> tmp is [N_obs × N_ens]
    tmp = (cov_gg + scaled_obs_noise_cov) \ (y - g)
    u += (cov_ug * tmp) # [N_par × N_ens]

    # store new parameters (and model outputs)
    push!(ekp.u, DataContainer(u, data_are_columns=true))
    push!(ekp.g, DataContainer(g, data_are_columns=true))

    # Store error
    compute_error!(ekp)

    # Check convergence
    cov_new = cov(get_u_final(ekp), dims=2)
    cov_ratio = det(cov_new) / det(cov_init)
    if cov_ratio < cov_threshold
        @warn string("New ensemble covariance determinant is less than ",
                     cov_threshold, " times its previous value.",
                     "\nConsider reducing the EK time step.")
    end
end

function update_ensemble!(ekp::MiniBatchKalmanProcess{FT, IT, Inversion},
                          g::Array{FT,2}, obs_mean_in::Array{FT}, obs_cov_in::Array{FT, 2};
                          cov_threshold::FT=0.01,
                          Δt_new=nothing,
                          deterministic_forward_map=true,) where {FT, IT}

    # Update follows eqns. (4) and (5) of Schillings and Stuart (2017)

    #catch works when g non-square
    if !(size(g)[2] == ekp.N_ens)
         throw(DimensionMismatch("ensemble size in EnsembleKalmanProcess and g do not match, try transposing g or check ensemble size"))
    end
    update_observations!(ekp, obs_mean_in, obs_cov_in)

    # u: N_par × N_ens
    # g: N_obs × N_ens
    u = get_u_final(ekp)
    N_obs = size(g, 1)

    cov_init = cov(u, dims=2)
    cov_ug = cov(u,g, dims = 2, corrected=false) # [N_par × N_obs]
    cov_gg = cov(g,g, dims = 2, corrected=false) # [N_obs × N_obs]

    if !isnothing(Δt_new)
        push!(ekp.Δt, Δt_new)
    elseif isnothing(Δt_new) && isempty(ekp.Δt)
        push!(ekp.Δt, FT(1))
    else
        push!(ekp.Δt, ekp.Δt[end])
    end

    # Scale noise using Δt
    scaled_obs_noise_cov = get_obs_cov_final(ekp) / ekp.Δt[end]
    noise = rand(MvNormal(zeros(N_obs), scaled_obs_noise_cov), ekp.N_ens)

    # Add obs_mean (N_obs) to each column of noise (N_obs × N_ens) if
    # G is deterministic
    y = deterministic_forward_map ? (get_obs_final(ekp) .+ noise) : (get_obs_final(ekp) .+ zero(noise))

    # N_obs × N_obs \ [N_obs × N_ens]
    # --> tmp is [N_obs × N_ens]
    tmp = (cov_gg + scaled_obs_noise_cov) \ (y - g)
    u += (cov_ug * tmp) # [N_par × N_ens]

    # store new parameters (and model outputs)
    push!(ekp.u, DataContainer(u, data_are_columns=true))
    push!(ekp.g, DataContainer(g, data_are_columns=true))

    # Store error
    compute_error!(ekp)

    # Check convergence
    cov_new = cov(get_u_final(ekp), dims=2)
    cov_ratio = det(cov_new) / det(cov_init)
    if cov_ratio < cov_threshold
        @warn string("New ensemble covariance determinant is less than ",
                     cov_threshold, " times its previous value.",
                     "\nConsider reducing the EK time step.")
    end
end

function update_ensemble!(ekp::EnsembleKalmanProcess{FT, IT, Sampler{FT}}, g_in::Array{FT,2}) where {FT, IT}

    #catch works when g_in non-square
    if !(size(g_in)[2] == ekp.N_ens)
         throw(DimensionMismatch("ensemble size in EnsembleKalmanProcess and g_in do not match, try transposing or check ensemble size"))
    end

    # u: N_ens × N_par
    # g: N_ens × N_obs
    u_old = get_u_final(ekp)
    u_old = permutedims(u_old,(2,1))
    u = u_old
    g = permutedims(g_in, (2,1))

    # u_mean: N_par × 1
    u_mean = mean(u', dims=2)
    # g_mean: N_obs × 1
    g_mean = mean(g', dims=2)
    # g_cov: N_obs × N_obs
    g_cov = cov(g, corrected=false)
    # u_cov: N_par × N_par
    u_cov = cov(u, corrected=false)

    # Building tmp matrices for EKS update:
    E = g' .- g_mean
    R = g' .- ekp.obs_mean
    # D: N_ens × N_ens
    D = (1/ekp.N_ens) * (E' * (ekp.obs_noise_cov \ R))

    Δt = 1/(norm(D) + 1e-8)

    noise = MvNormal(u_cov)

    implicit = (1 * Matrix(I, size(u)[2], size(u)[2]) + Δt * (ekp.process.prior_cov' \ u_cov')') \
                  (u'
                    .- Δt * ( u' .- u_mean) * D
                    .+ Δt * u_cov * (ekp.process.prior_cov \ ekp.process.prior_mean)
                  )

    u = implicit' + sqrt(2*Δt) * rand(noise, ekp.N_ens)'

    # store new parameters (and model outputs)
    push!(ekp.u, DataContainer(u, data_are_columns=false))
    push!(ekp.g, DataContainer(g, data_are_columns=false))
    # u_old is N_ens × N_par, g is N_ens × N_obs,
    # but stored in data container with N_ens as the 2nd dim

    compute_error!(ekp)

end

# outer constructors
function EnsembleKalmanProcess(
    obs_mean::Array{FT, 1},
    obs_noise_cov::Array{FT, 2},
    process::Unscented{FT, IT};
    Δt=FT(1)) where {FT<:AbstractFloat, IT<:Int}

    #initial parameters stored as columns
    init_params = [DataContainer(update_ensemble_prediction!(process), data_are_columns=true)]
    # Number of parameter
    N_u = length(process.u_mean[1])
    # ensemble size
    N_ens = 2*N_u + 1 #stored with data as columns
    #store for model evaluations
    g=[]
    # error store
    err = FT[]
    # timestep store
    Δt = Array([Δt])

    EnsembleKalmanProcess{FT, IT, Unscented}(init_params, obs_mean, obs_noise_cov, N_ens, g, err, Δt, process)
end

function MiniBatchKalmanProcess(
    process::Unscented{FT, IT};
    Δt=FT(1)) where {FT<:AbstractFloat, IT<:Int}

    #initial parameters stored as columns
    init_params = [DataContainer(update_ensemble_prediction!(process), data_are_columns=true)]
    # Number of parameter
    N_u = length(process.u_mean[1])
    # ensemble size
    N_ens = 2*N_u + 1 #stored with data as columns
    #store for model evaluations
    g=[]
    # store for observations
    obs_mean = []
    obs_noise_cov = []
    # error store
    err = FT[]
    # timestep store
    Δt = Array([Δt])

    MiniBatchKalmanProcess{FT, IT, Unscented}(init_params, obs_mean, obs_noise_cov, N_ens, g, err, Δt, process)
end

"""
EnsembleKalmanProcess Constructor
u0_mean::Array{FT} : prior mean
uu0_cov::Array{FT, 2} : prior covariance
obs_mean::Array{FT,1} : observation
obs_noise_cov::Array{FT, 2} : observation error covariance
α_reg::FT : regularization parameter toward u0 (0 < α_reg <= 1), default should be 1, without regulariazion
update_freq::IT : set to 0 when the inverse problem is not identifiable,
                  namely the inverse problem has multiple solutions,
                  the covariance matrix will represent only the sensitivity of the parameters,
                  instead of posterior covariance information;
                  set to 1 (or anything > 0) when the inverse problem is identifiable, and
                  the covariance matrix will converge to a good approximation of the
                  posterior covariance with an uninformative prior.

"""
function Unscented(
    u0_mean::Array{FT, 1},
    uu0_cov::Array{FT, 2},
    α_reg::FT,
    update_freq::IT;
    modified_unscented_transform::Bool = true,
    κ::FT = 0.0,
    β::FT = 2.0) where {FT<:AbstractFloat, IT<:Int}

    N_u = size(u0_mean,1)
    # ensemble size
    N_ens = 2*N_u+1

    c_weights = zeros(FT, N_u)
    mean_weights = zeros(FT, N_ens)
    cov_weights = zeros(FT, N_ens)

    # todo parameters λ, α, β

    α = min(sqrt(4/(N_u + κ)), 1.0)
    λ = α^2*(N_u + κ) - N_u

    c_weights[1:N_u]     .=  sqrt(N_u + λ)
    mean_weights[1] = λ/(N_u + λ)
    mean_weights[2:N_ens] .= 1/(2(N_u + λ))
    cov_weights[1] = λ/(N_u + λ) + 1 - α^2 + β
    cov_weights[2:N_ens] .= 1/(2(N_u + λ))

    if modified_unscented_transform
        mean_weights[1] = 1.0
        mean_weights[2:N_ens] .= 0.0
    end

    u_mean = Array{FT,1}[]  # array of Array{FT, 2}'s
    push!(u_mean, u0_mean) # insert parameters at end of array (in this case just 1st entry)
    uu_cov = Array{FT,2}[] # array of Array{FT, 2}'s
    push!(uu_cov, uu0_cov) # insert parameters at end of array (in this case just 1st entry)

    obs_pred = Array{FT, 1}[]  # array of Array{FT, 2}'s
    err = FT[]

    Σ_ω = (2 - α_reg^2)*uu0_cov
    Σ_ν_scale = 2.0

    r = u0_mean
    iter = 0

    Unscented(u_mean, uu_cov,  obs_pred, c_weights, mean_weights, cov_weights, Σ_ω, Σ_ν_scale, α_reg, r, update_freq, iter)
end

"""
construct_sigma_ensemble
Construct the sigma ensemble, based on the mean x_mean, and covariance x_cov
"""
function construct_sigma_ensemble(process::Unscented, x_mean::Array{FT}, x_cov::Array{FT,2}) where {FT<:AbstractFloat, IT<:Int}

    N_x = size(x_mean,1)

    c_weights = process.c_weights

    chol_xx_cov = cholesky(Hermitian(x_cov)).L

    x = zeros(FT, N_x, 2*N_x+1)
    x[:, 1] = x_mean
    for i = 1: N_x
        x[:, i+1] = x_mean + c_weights[i]*chol_xx_cov[:,i]
        x[:, i+1+N_x] = x_mean - c_weights[i]*chol_xx_cov[:,i]
    end

    return x
end

"""
construct_mean x_mean from ensemble x
"""
function construct_mean(uki::Union{EnsembleKalmanProcess{FT, IT,Unscented}, MiniBatchKalmanProcess{FT, IT,Unscented}},
    x::Array{FT,2}) where {FT<:AbstractFloat, IT<:Int}
    N_x, N_ens = size(x)

    @assert(uki.N_ens == N_ens)

    x_mean = zeros(FT, N_x)

    mean_weights = uki.process.mean_weights

    for i = 1: N_ens
        x_mean += mean_weights[i]*x[:, i]
    end

    return x_mean
end

"""
construct_cov xx_cov from ensemble x and mean x_mean
"""
function construct_cov(uki::Union{EnsembleKalmanProcess{FT, IT,Unscented}, MiniBatchKalmanProcess{FT, IT,Unscented}},
    x::Array{FT,2}, x_mean::Array{FT}) where {FT<:AbstractFloat, IT<:Int}
    N_ens, N_x = uki.N_ens, size(x_mean,1)

    cov_weights = uki.process.cov_weights

    xx_cov = zeros(FT, N_x, N_x)

    for i = 1: N_ens
        xx_cov .+= cov_weights[i]*(x[:,i] - x_mean)*(x[:,i] - x_mean)'
    end

    return xx_cov
end

"""
construct_cov xy_cov from ensemble x and mean x_mean, ensemble obs_mean and mean y_mean
"""
function construct_cov(uki::Union{EnsembleKalmanProcess{FT, IT,Unscented}, MiniBatchKalmanProcess{FT, IT,Unscented}}, x::Array{FT,2}, x_mean::Array{FT}, obs_mean::Array{FT,2}, y_mean::Array{FT}) where {FT<:AbstractFloat, IT<:Int, P<:Process}
    N_ens, N_x, N_y = uki.N_ens, size(x_mean,1), size(y_mean,1)

    cov_weights = uki.process.cov_weights

    xy_cov = zeros(FT, N_x, N_y)

    for i = 1: N_ens
        xy_cov .+= cov_weights[i]*(x[:, i] - x_mean)*(obs_mean[:, i] - y_mean)'
    end

    return xy_cov
end

"""
uki prediction step : generate sigma points
"""
function update_ensemble_prediction!(process::Unscented; Δt_new=1.0) where {FT<:AbstractFloat, IT<:Int}

    process.iter += 1
    # update evolution covariance matrix
    if process.update_freq > 0 && process.iter%process.update_freq == 0
        process.Σ_ω = (2 - process.α_reg^2)process.uu_cov[end]
    end

    u_mean  = process.u_mean[end]
    uu_cov = process.uu_cov[end]

    α_reg = process.α_reg
    r = process.r
    Σ_ω = process.Σ_ω * Δt_new

    N_u = length(process.u_mean[1])
    ############# Prediction step:

    u_p_mean  = α_reg*u_mean + (1-α_reg)*r
    uu_p_cov = α_reg^2*uu_cov + Σ_ω

    ############ Generate sigma points
    u_p = construct_sigma_ensemble(process, u_p_mean, uu_p_cov)
    return u_p
end

"""
uki analysis step
g is the predicted observations  Ny  by N_ens matrix
"""
function update_ensemble_analysis!(uki::EnsembleKalmanProcess{FT, IT,Unscented}, u_p::Array{FT, 2}, g::Array{FT, 2}; Δt_new = 1) where {FT<:AbstractFloat, IT<:Int}

    obs_mean = uki.obs_mean
    Σ_ν = uki.process.Σ_ν_scale * uki.obs_noise_cov / Δt_new

    N_u, N_y, N_ens = length(uki.process.u_mean[1]), length(uki.obs_mean), uki.N_ens
    ############# Prediction step:

    u_p_mean = construct_mean(uki, u_p)
    uu_p_cov = construct_cov(uki, u_p, u_p_mean)

    ###########  Analysis step

    g_mean = construct_mean(uki, g)
    gg_cov = construct_cov(uki, g, g_mean) + Σ_ν
    ug_cov = construct_cov(uki, u_p, u_p_mean, g, g_mean)

    tmp = ug_cov/gg_cov

    u_mean =  u_p_mean + tmp*(obs_mean - g_mean)
    uu_cov =  uu_p_cov - tmp*ug_cov'

    ########### Save results
    push!(uki.process.obs_pred, g_mean) # N_ens x N_data
    push!(uki.process.u_mean, u_mean) # N_ens x N_params
    push!(uki.process.uu_cov, uu_cov) # N_ens x N_data

    push!(uki.g, DataContainer(g, data_are_columns=true))

    if !isnothing(Δt_new)
        push!(uki.Δt, Δt_new)
    elseif isnothing(Δt_new) && isempty(uki.Δt)
        push!(uki.Δt, FT(1))
    else
        push!(uki.Δt, ekp.Δt[end])
    end

    compute_error!(uki)
end

"""
uki analysis step
g is the predicted observations  Ny  by N_ens matrix
"""
function update_ensemble_analysis!(uki::MiniBatchKalmanProcess{FT, IT,Unscented}, u_p::Array{FT, 2}, g::Array{FT, 2}; Δt_new = 1) where {FT<:AbstractFloat, IT<:Int}

    obs_mean = get_obs_final(uki)
    Σ_ν = uki.process.Σ_ν_scale * get_obs_cov_final(uki) / Δt_new

    N_u, N_y, N_ens = length(uki.process.u_mean[1]), length(obs_mean), uki.N_ens
    ############# Prediction step:

    u_p_mean = construct_mean(uki, u_p)
    uu_p_cov = construct_cov(uki, u_p, u_p_mean)

    ###########  Analysis step

    g_mean = construct_mean(uki, g)
    gg_cov = construct_cov(uki, g, g_mean) + Σ_ν
    ug_cov = construct_cov(uki, u_p, u_p_mean, g, g_mean)

    tmp = ug_cov/gg_cov

    u_mean =  u_p_mean + tmp*(obs_mean - g_mean)
    uu_cov =  uu_p_cov - tmp*ug_cov'

    ########### Save results
    push!(uki.process.obs_pred, g_mean) # N_ens x N_data
    push!(uki.process.u_mean, u_mean) # N_ens x N_params
    push!(uki.process.uu_cov, uu_cov) # N_ens x N_data

    push!(uki.g, DataContainer(g, data_are_columns=true))

    if !isnothing(Δt_new)
        push!(uki.Δt, Δt_new)
    elseif isnothing(Δt_new) && isempty(uki.Δt)
        push!(uki.Δt, FT(1))
    else
        push!(uki.Δt, ekp.Δt[end])
    end

    compute_error!(uki)
end

function update_ensemble!(uki::EnsembleKalmanProcess{FT, IT, Unscented}, g_in::Array{FT, 2}; Δt_new = 1) where {FT<:AbstractFloat, IT<:Int}
    #catch works when g_in non-square
    if !(size(g_in)[2] == uki.N_ens)
        throw(DimensionMismatch("ensemble size in EnsembleKalmanProcess and g_in do not match, try transposing or check ensemble size"))
    end

    u_p_old = get_u_final(uki)

    #perform analysis on the model runs
    update_ensemble_analysis!(uki, u_p_old, g_in, Δt_new=Δt_new)
    #perform new prediction output to model parameters u_p
    u_p = update_ensemble_prediction!(uki.process, Δt_new=Δt_new)

    push!(uki.u, DataContainer(u_p, data_are_columns=true))

    return u_p

end

function update_ensemble!(uki::MiniBatchKalmanProcess{FT, IT, Unscented}, g_in::Array{FT, 2},
        obs_mean_in::Array{FT}, obs_cov_in::Array{FT, 2}; Δt_new = 1) where {FT<:AbstractFloat, IT<:Int}
    #catch works when g_in non-square
    if !(size(g_in)[2] == uki.N_ens)
        throw(DimensionMismatch("ensemble size in EnsembleKalmanProcess and g_in do not match, try transposing or check ensemble size"))
    end

    u_p_old = get_u_final(uki)

    #perform analysis on the model runs
    update_observations!(uki, obs_mean_in, obs_cov_in)
    update_ensemble_analysis!(uki, u_p_old, g_in, Δt_new=Δt_new)
    #perform new prediction output to model parameters u_p
    u_p = update_ensemble_prediction!(uki.process, Δt_new=Δt_new)

    push!(uki.u, DataContainer(u_p, data_are_columns=true))

    return u_p

end

function update_observations!(ekp::MiniBatchKalmanProcess, obs_mean_in::Array{FT}, obs_cov_in::Array{FT, 2}) where {FT<:AbstractFloat}

    push!(ekp.obs_mean, obs_mean_in)
    push!(ekp.obs_noise_cov, DataContainer(obs_cov_in, data_are_columns=true))

    return
end

function get_u_mean_final(uki::Union{EnsembleKalmanProcess{FT, IT,Unscented}, MiniBatchKalmanProcess{FT, IT,Unscented}}) where {FT<:AbstractFloat, IT<:Int}
    return uki.process.u_mean[end]
end

function compute_error!(uki::EnsembleKalmanProcess{FT, IT,Unscented}) where {FT<:AbstractFloat, IT<:Int}
    mean_g = uki.process.obs_pred[end]
    diff = uki.obs_mean - mean_g
    X = uki.obs_noise_cov \ diff # diff: column vector
    newerr = dot(diff, X)
    push!(uki.err, newerr)
end

function compute_error!(uki::MiniBatchKalmanProcess{FT, IT,Unscented}) where {FT<:AbstractFloat, IT<:Int}
    mean_g = uki.process.obs_pred[end]
    diff = get_obs_final(uki) - mean_g
    X = get_obs_cov_final(uki) \ diff # diff: column vector
    newerr = dot(diff, X)
    push!(uki.err, newerr)
end

function Gaussian_2d(u_mean::Array{FT,1}, uu_cov::Array{FT,2}, Nx::IT, Ny::IT; xx = nothing, yy = nothing) where {FT<:AbstractFloat, IT<:Int}
    # 2d Gaussian plot
    u_range = [min(5*sqrt(uu_cov[1,1]), 5); min(5*sqrt(uu_cov[2,2]), 5)]

    if xx === nothing
        xx = Array(LinRange(u_mean[1] - u_range[1], u_mean[1] + u_range[1], Nx))
    end
    if yy == nothing
        yy = Array(LinRange(u_mean[2] - u_range[2], u_mean[2] + u_range[2], Ny))
    end
    X,Y = repeat(xx, 1, Ny), Array(repeat(yy, 1, Nx)')
    Z = zeros(FT, Nx, Ny)

    det_uu_cov = det(uu_cov)

    for ix = 1:Nx
        for iy = 1:Ny
            Δxy = [xx[ix] - u_mean[1]; yy[iy] - u_mean[2]]
            Z[ix, iy] = exp(-0.5*(Δxy'/uu_cov*Δxy)) / (2 * pi * sqrt(det_uu_cov))
        end
    end

    return xx, yy, Z

end
end # module


text:
# Note that this script can accept some limited command-line arguments, run
# `julia build_tarballs.jl --help` to see a usage message.
using BinaryBuilder

name = "obstack"
version = v"1.1.0"

# Collection of sources required to build obstack
sources = [
    ArchiveSource("https://github.com/pullmoll/musl-obstack/archive/v$(version.major).$(version.minor).tar.gz",
                  "52a216613e7d55e8725e43d017bb2d49a4b1ffa1e06da472f03c7f9875df7d0d"),
]

# Bash recipe for building across all platforms
script = raw"""
cd $WORKSPACE/srcdir/musl-obstack-*/
./bootstrap.sh
CFLAGS="-fPIC" ./configure --prefix=${prefix} --build=${MACHTYPE} --host=${target}
make -j${nproc}
make install
"""

# Select Unix platforms
platforms = [p for p in supported_platforms() if Sys.islinux(p) && libc(p) == "musl"]

# The products that we will ensure are always built
products = [
    LibraryProduct("libobstack", :libobstack),
]

# Dependencies that must be installed before this package can be built
dependencies = Dependency[
]

# Build the tarballs, and possibly a `build.jl` as well.
build_tarballs(ARGS, name, version, sources, script, platforms, products, dependencies)

max_stars_repo_name: bzinberg/Yggdrasil

max_issues_repo_name: bzinberg/Yggdrasil

max_forks_repo_name: bzinberg/Yggdrasil


text:
# module Battlesnake

# export Boardstate, Snake, Position, Direction, up, down, left, right, move_from

# ============== DIRECTION ==================#
@enum Direction up down left right

# ============== POSITION ==================#
struct Position
    x::Int
    y::Int
end

function get_neighbors(pos::Position)
    return [(dir, move_from(pos, dir)) for dir in [up, down, left, right]]
end

function move_from(pos::Position, dir::Direction)
    if dir == up
        return Position(pos.x, pos.y + 1)
    elseif dir == down
        return Position(pos.x, pos.y - 1)
    elseif dir == right
        return Position(pos.x + 1, pos.y)
    elseif dir == left
        return Position(pos.x - 1, pos.y)
    end
end

# ============== FOOD ==================#
mutable struct Food
    positions::Vector{Position}
end

# ============== SNAKE ==================#

mutable struct Snake
    body::Vector{Position}
    ate_food::Bool
    health::Int
end

function get_head(snake::Snake)
    return snake.body[1]
end

function get_body_no_head(snake::Snake)
    return snake.body[2:length(snake.body)]
end

function possible_actions(snake::Snake)
    head = snake.body[1]
    neck = snake.body[2]
    # coming from right
    if head.x < neck.x
        return [up, down, left]
    end
    # coming from left
    if head.x > neck.x
        return [up, down, right]
    end

    # coming from below
    if head.y > neck.y
        return [up, right, left]
    end

    # coming from above
    if head.y < neck.y
        return [down, right, left]
    end
    return [up, down, left, right]
end

function move_snake!(snake::Snake, food::Food, action::Direction)
    h = get_head(snake)
    new_head = move_from(h, action)
    insert!(snake.body, 1, new_head)
    if !snake.ate_food
        pop!(snake.body)
    end
    snake.ate_food = false
    snake.health -= 1
    food_eaten = -1
    for i in 1:length(food.positions)
        if new_head == food.positions[i]
            snake.ate_food = true
            food_eaten = i
            snake.health = 100
        end
    end
    if food_eaten > -1
        deleteat!(food.positions, food_eaten)
    end
end

function check_bounds(pos::Position, height::Int, width::Int)
    return pos.x >= 1 && pos.x < width + 1 && pos.y >= 1 && pos.y < height + 1
end

function check_dead(snakes::Vector{Snake}, height::Int, width::Int)
    dead = Set()
    for (p, snek) in enumerate(snakes)
        # starve
        if snek.health == 0
            push!(dead, p)
        end

        for (o, opponent) in enumerate(snakes)
            if o == p
                continue
            end

            if get_head(snek) in opponent.body
                # head2head
                if get_head(snek) == get_head(opponent)
                    # RAM
                    if length(snek.body) > length(opponent.body)
                        push!(dead, o)
                    # draw
                    elseif length(snek.body) == length(opponent.body)
                        push!(dead, o)
                        push!(dead, p)
                    else
                        # lost (head2head) collision with opponent
                        push!(dead, p)
                    end
                else
                    # ran into opponent
                    push!(dead, p)
                end
            end

            # collision with self
            if get_head(snek) in get_body_no_head(snek)
                push!(dead, p)
            end
            if !check_bounds(get_head(snek), height, width)
                push!(dead, p)
            end

        end
    end
    return dead
end

# ============== BOARDSTATE ==================#
mutable struct Boardstate
    snakes::Vector{Snake}
    food::Food
    height::Int
    width::Int
end

function is_occupied_by_snake(pos::Position, state::Boardstate)
    for snake in state.snakes
        if pos in snake.body
            return true
        end
    end
    return false
end

function print_board(state::Boardstate)
    board = fill(0, (state.height, state.width))
    for (i,snek) in enumerate(state.snakes)
        for p in snek.body
            board[p.y, p.x] = i
        end
    end
    for f in state.food.positions
        board[f.y, f.x] = 5
    end
    reverse!(board, dims = 1)
    for r in 1:size(board,1)
        println(board[r, :])
    end
end

function next_player(current_player, num_players)
    return (current_player % num_players) + 1
end

# end # module end

max_stars_repo_name: rhotertj/solide-schlange

max_issues_repo_name: rhotertj/solide-schlange

max_forks_repo_name: rhotertj/solide-schlange