text:
# ====================================================================
# Component structure
#
mutable struct Lorentzian_1D <: AbstractComponent
    norm::Parameter
    center::Parameter
    fwhm::Parameter

    function Lorentzian_1D(norm::Number, center::Number, fwhm::Number)
        @assert norm  > 0
        @assert fwhm > 0

        out = new(Parameter(norm), Parameter(center), Parameter(fwhm))
        out.norm.low = 0
        out.fwhm.low = 0
        return out
    end
end

mutable struct Lorentzian_2D <: AbstractComponent
    norm::Parameter
    centerX::Parameter
    centerY::Parameter
    fwhmX::Parameter
    fwhmY::Parameter

    function Lorentzian_2D(norm::Number, centerX::Number, centerY::Number, fwhmX::Number, fwhmY::Number)
        @assert norm   > 0
        @assert fwhmX > 0
        @assert fwhmY > 0

        out = new(Parameter(norm), Parameter(centerX), Parameter(centerY), Parameter(fwhmX), Parameter(fwhmY))
        out.norm.low = 0
        out.fwhmX.low = 0
        out.fwhmY.low = 0

        return out
    end
end

Lorentzian(norm, center, fwhm) = Lorentzian_1D(norm, center, fwhm)
Lorentzian(norm, centerX, centerY, fwhmX, fwhmY) = Lorentzian_2D(norm, centerX, centerY, fwhmX, fwhmY)

# ====================================================================
# Evaluate component
function evaluate!(buffer, comp::Lorentzian_1D, x::AbstractDomain{1},
                   norm, center, fwhm)
    @. (buffer = norm /
        (1. +
         ((x - center) / fwhm)^2.
        ))
end

function evaluate!(buffer, comp::Lorentzian_2D, domain::AbstractDomain{2},
                   norm, centerX, centerY, fwhmX, fwhmY)
    x = domain[1]
    y = domain[2]

    @. (buffer = norm /
        (1. +
         ((x - centerX) / fwhmX)^2. +
         ((y - centerY) / fwhmY)^2.
         ))
end


text:
function sortall(xs::AbstractArray...; kw...)
    p = sortperm(first(xs); kw...)
    map(x -> x[p], xs)
end

function sortall!(xs::AbstractArray...; kw...)
    p = sortperm(first(xs); kw...)
    for x in xs
        permute!(x, p)
    end
    return xs
end

function write_feaimpt(w, c; dst = "feaimpt.csv")
    length(w) != length(c) && return
    w, c = sortall(vec(w), vec(c), by = abs, rev = true)
    open(dst, "w") do io
        for i in 1:min(100, length(c))
            write(io, togbk(c[i]), ',')
            println(io, trunc(w[i], digits = 2))
        end
    end
end

function hasgpu()
    try
        run(`nvidia-smi`, wait = false)
        return true
    catch
        return false
    end
end

usegpu() = hasgpu() && get(ENV, "USE_GPU", "1") == "1"

countgpus() = hasgpu() ? count(x -> x == '\n', read(`nvidia-smi -L`, String)) : 0

function to_svm(x, y, w = nothing; group = nothing)
    @imports numpy as np
    if !isnothing(group) && length(group) != length(y)
        qid = np.repeat(1:length(group), group)
    else
        qid = group
    end
    dst = joinpath("/dev/shm", tempname() * ".svmlight")
    x = reshape(x, size(x, 1), :)
    y, qid = vec(y), vec(qid)
    dump_svmlight_file(x, y, dst, query_id = qid, zero_based = false)
    return dst
end

function to_svm(x)
    x = reshape(x, size(x, 1), :)
    y = zeros(Float32, size(x, 2))
    to_svm(x, y, group = [length(y)])
end

macro NT(xs...)
    xs = [:($x = $x) for x in xs]
    esc(:(($(xs...),)))
end

function bash(str, exe = run)
    exe(`bash -c $str`)
end

parseenv(key, default::String) = get(ENV, string(key), string(default))

function parseenv(key, default::T) where T
    str = get(ENV, string(key), string(default))
    if hasmethod(parse, (Type{T}, String))
        parse(T, str)
    else
        include_string(Main, str)
    end
end

rmdir(src) = rm(src, recursive = true)

memory(x) = Base.summarysize(x) / 1024^2

function align(x, y)
    iy, Nx, Ny = 1, length(x), length(y)
    for ix in 1:(Nx - Ny + 1)
        for iy in 1:Ny
            if x[iy + ix - 1] != y[iy]
                break
            elseif iy == Ny
                return ix
            end
        end
    end
    return 0
end

macro gc(exs...)
    Expr(:block, [:($ex = 0) for ex in exs]..., :(@eval GC.gc(true))) |> esc
end

macro redirect(src, ex)
    src = src == :devnull ? "/dev/null" : src
    quote
        io = open($(esc(src)), "a")
        o, e = stdout, stderr
        redirect_stdout(io)
        redirect_stderr(io)
        try
            $(esc(ex)); sleep(0.01)
        finally
            flush(io); close(io)
            redirect_stdout(o)
            redirect_stderr(e)
        end
    end
end

Base.vec(::Nothing) = nothing

bra(x) = isnothing(x) ? x : [x]

sigmoid(x) = 1 / (1 + exp(-x))

function softmax(xs::AbstractArray; dims=1)
    max_ = maximum(xs, dims = dims)
    exp_ = exp.(xs .- max_)
    exp_ ./ sum(exp_, dims = dims)
end

export @trys
macro trys(exs...)
    expr = :()
    for ex in exs[end:-1:1]
        expr = :(try $ex; catch e; $expr; end)
    end
    esc(expr)
end

ext: jl


text:
# Holds the possible QP-like objective structs

#############################################################
###  Structs for Quadratic and Linear Objective Functions ###
#############################################################

abstract type objectiveFunc end
abstract type objectiveQP_abstract <: objectiveFunc end

using SparseArrays

@doc raw"""
    objectiveQP(Q, p)

Quadratic (and Linear) Objective/Cost Function
----

```math
\frac{1}{2} x^{⊤} Q x + p^{⊤} x
```

`Q` is an nxn Symmetric Matrix (that is positive definite for convex)

`p` is an nx1 vector

See also: [`objectiveQP`](@ref), [`fObjQP`](@ref), [`dfdxQP`](@ref)
"""
struct objectiveQP <: objectiveQP_abstract
    "Positive Definite Symmetric Matrix"
    Q
    p
end

@doc raw"""
    objectiveQ(Q)

Quadratic (only) Objective/Cost Function
----

```math
\frac{1}{2} x^{⊤} Q x
```

`Q` is an nxn Symmetric Matrix (that is positive definite for convex)

See also: [`objectiveQP`](@ref) [`fObjQP`](@ref), [`dfdxQP`](@ref)
"""
struct objectiveQ <: objectiveQP_abstract
    "Positive Definite Symmetric Matrix"
    Q
end

@doc raw"""
    objectiveP(P)

Linear (only) Objective/Cost Function
----

```math
p^{⊤} x
```

`p` is an nx1 vector

See also: [`objectiveQP`](@ref), [`fObjQ`](@ref), [`dfdxQ`](@ref)
"""
struct objectiveP <: objectiveQP_abstract
    "Positive Definite Symmetric Matrix"
    p
end

#############################################################
###       Evaluate Quadratic and Linear Objectives        ###
#############################################################

"""
    fObjQP(qp::objectiveQP, x)

Evaluates a quadratic function at the input `x`. The input `x` should be a
*column* vector (i.e. `x = [4; 3]`)
"""
function fObjQP(qp::objectiveQP, x)
    return (1/2) * x' * (qp.Q) * x + (qp.p)' * x
end

"""
    fObjQP(q::objectiveQ, x)

Evaluates a quadratic (only) function at the input `x`. The input `x` should
be a *column* vector (i.e. `x = [4; 3]`)
"""
function fObjQP(q::objectiveQ, x)
    return (1/2) * x' * (q.Q) * x
end

"""
    fObjQP(p::objectiveP, x)

Evaluates a linear function at the input `x`. The input `x` should be a
*column* vector (i.e. `x = [4; 3]`)
"""
function fObjQP(p::objectiveP, x)
    return (p.p)' * x
end

#############################################################
###     Gradients for Quadratic and Linear Objectives     ###
#############################################################

@doc raw"""
    dfdxQP(qp::objectiveQP, x)

Evaluates the derivative of a quadratic function at input `x`. The input `x`
should be a *column* vector (i.e. `x = [4; 3]`)

```math
\frac{d}{dx} (\frac{1}{2} x^{⊤} Q x + p^{⊤} x) = Qx + p
```
"""
function dfdxQP(qp::objectiveQP, x)
    return (qp.Q) * x + qp.p
end

@doc raw"""
    dfdxQ(q::objectiveQ, x)

Evaluates the derivative of a quadratic (only) function at input `x`. The
input `x` should be a *column* vector (i.e. `x = [4; 3]`)

```math
\frac{d}{dx} (\frac{1}{2} x^{⊤} Q x) = Qx
```
"""
function dfdxQP(q::objectiveQ, x)
    return (q.Q) * x
end

@doc raw"""
    dfdxQP(p::objectiveP, x = 0)

Evaluates the derivative of a linear function. The `x` is unused, but is kept
to maintain the same function header.

```math
\frac{d}{dx} (p^{⊤} x) = p
```
"""
function dfdxQP(p::objectiveP, x = 0)
    return p.p
end

#############################################################
###      Hessians for Quadratic and Linear Objectives     ###
#############################################################

@doc raw"""
    hessQP(qp::objectiveQP)

Evaluates the hessian of a quadratic function.

```math
H(\frac{1}{2} x^{⊤} Q x + p^{⊤} x) = Q
```

Note that Q is already square, symmetric, and (generally) positive definite.
"""
function hessQP(qp::objectiveQP)
    return qp.Q
end

@doc raw"""
    hessQP(q::objectiveQ)

Evaluates the hessian of a quadratic (only) function.

```math
H(\frac{1}{2} x^{⊤} Q x) = Q
```

Note that Q is already square, symmetric, and (generally) positive definite.
"""
function hessQP(q::objectiveQ)
    return q.Q
end

@doc raw"""
    hessQP(q::objectiveQ)

Evaluates the hessian of a linear (only) function. (which is zero)

```math
H(p^{⊤} x) = 0
```
"""
function hessQP(p::objectiveP)
    n = size(p.p, 1)
    return spzeros(n, n)
end

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

max_stars_repo_name: danineamati/TrajOptSOCPs


text:
module MomentBasedBayesianEstimators
using Reexport
@reexport using MomentBasedEstimators
@reexport using Divergences
import MomentBasedEstimators: MinimumDivergenceProblem

immutable BayesianMomentModel{F, R <: MinimumDivergenceProblem}
    f!::Function
    prior!::Function
    data::F
    m::R
    p::Array{Float64, 1}
end

flatprior!(a, theta) = return a

function BayesianMomentModel(f!::Function, prior!::Function, m::MinimumDivergenceProblem, data)
    p = zeros(1)
    BayesianMomentModel(f!, prior!, data, m, p)
end

function posterior(b::BayesianMomentModel, theta::Vector{Float64})
    ## f! is a mutating function
    b.f!(b.m.e.mm.S, theta, b.data)
    solve!(b.m)
    sum(log(b.m.m.inner.x::Array{Float64, 1})) + b.prior!(b.p, theta)
end

export BayesianMomentModel, flatprior!, posterior

# package code goes here

end # module


text:
# Author: Mathias Louboutin, mlouboutin3@gatech.edu
# Updated July 2020

ftol = 1f-5

@testset "PhysicalParameter Unit Tests in $(nd) dimensions" for nd=[2, 3]
    n = Tuple(11+i for i=1:nd)
    d = Tuple(1f0+i for i=1:nd)
    o = Tuple(0f0+i for i=1:nd)
    a = randn(Float32, n...)
    p = PhysicalParameter(a, d, o)
    @test PhysicalParameter(p) == p
    @test PhysicalParameter(p, n, d, o) == p
    @test PhysicalParameter(p.data, n, d, o) == p
    @test PhysicalParameter(vec(p.data), n, d, o) == p

    @test size(p) == (prod(n), 1)
    @test size(conj(p)) == (prod(n), 1)

    @test size(adjoint(p)) == (prod(n), 1)
    @test adjoint(p).n == p.n[end:-1:1]
    @test adjoint(p).d == p.d[end:-1:1]
    @test adjoint(p).o == p.o[end:-1:1]
    @test isapprox(adjoint(p).data, permutedims(p.data, nd:-1:1))

    @test size(transpose(p)) == (prod(n), 1)
    @test transpose(p).n == p.n[end:-1:1]
    @test transpose(p).d == p.d[end:-1:1]
    @test transpose(p).o == p.o[end:-1:1]
    @test isapprox(transpose(p).data, permutedims(p.data, nd:-1:1))

    ps = similar(p, Model(n, d, o, a))
    @test norm(ps) == 0
    @test ps.n == n
    @test ps.d == d
    @test ps.o == o

    @test isequal(p.data, a)
    p .= zeros(n...)
    p .= a
    @test isequal(p.data, a)
    @test p.d == d
    @test p.n == n
    @test p.o == o

    # indexing
    @test firstindex(p) == 1
    @test lastindex(p) == prod(n)
    @test p == p
    @test isapprox(p, a)
    @test isapprox(a, p)

    # Subsampling
    inds = [3:11 for i=1:nd]
    b = p[inds...]
    @test b.o == tuple([o+d*2 for (o,d)=zip(p.o,p.d)]...)
    @test b.n == tuple([9 for i=1:nd]...)
    @test b.d == p.d

    inds = [3:2:11 for i=1:nd]
    b = p[inds...]
    @test b.o == tuple([o+d*2 for (o,d)=zip(p.o,p.d)]...)
    @test b.n == tuple([5 for i=1:nd]...)
    @test b.d == tuple([d*2 for d=p.d]...)

    # copies
    p2 = similar(p)
    @test p2.n == p.n
    @test norm(p2) == 0

    p2 = copy(p)
    @test p2.n == p.n
    @test norm(p2) == norm(p)

    p2 = 0f0 .* p .+ 1f0
    @test norm(p2, 1) == prod(n)

    # Some basics
    pl = PhysicalParameter(n, d, o)
    @test norm(pl) == 0
    pl = PhysicalParameter(0, n, d, o)
    @test pl.data == 0
    @test isapprox(dot(p, p), norm(p)^2)
    @test isapprox(dot(ones(n), p), sum(p.data))
    @test isapprox(dot(p, ones(n)), sum(p.data))

    # broadcast multiplication
    u = p
    v = p2
    u_scale = deepcopy(p)
    v_scale = deepcopy(p2)
    c = ones(Float32, v.n)

    u_scale .*= 2f0
    @test isapprox(u_scale, 2f0 * u; rtol=ftol)
    v_scale .+= 2f0
    @test isapprox(v_scale, 2f0 .+ v; rtol=ftol)
    u_scale ./= 2f0
    @test isapprox(u_scale, u; rtol=ftol)
    u_scale .= 2f0 .* u_scale .+ v_scale
    @test isapprox(u_scale, 2f0 * u .+ 2f0 + v; rtol=ftol)
    u_scale .= u .+ v
    @test isapprox(u_scale, u + v)
    u_scale .= u .- v
    @test isapprox(u_scale, u - v)
    u_scale .= u .* v
    @test isapprox(u_scale.data[1], u.data[1].*v.data[1])
    u_scale .= u ./ v
    @test isapprox(u_scale.data[1], u.data[1]./v.data[1])

    @test isapprox(u.*c, u)
    @test isapprox(c.*u, u)
    @test isapprox(u./c, u)
    @test isapprox(c./u, 1 ./u)
    @test isapprox(c .+ u, 1 .+ u)
    @test isapprox(c .- u, 1 .- u)
    @test isapprox(u .+ c, 1 .+ u)
    @test isapprox(u .- c, u .- 1)

    # Arithmetic
    v = PhysicalParameter(randn(Float32, n...), d, o)
    w = PhysicalParameter(randn(Float32, n...), d, o)
    a = .5f0 + rand(Float32)
    b = .5f0 + rand(Float32)
    c2 = randn(size(v, 1), size(v, 1))

    @test isapprox(u + (v + w), (u + v) + w; rtol=ftol)
    @test isapprox(u + v, v + u; rtol=ftol)
    @test isapprox(u, u .+ 0; rtol=ftol)
    @test iszero(norm(u + u.*(-1)))
    @test isapprox(-u, (-1f0).*u; rtol=ftol)
    @test isapprox(a .* (b .* u), (a * b) .* u; rtol=ftol)
    @test isapprox(u, u .* 1; rtol=ftol)
    @test isapprox(a .* (u + v), a .* u + a .* v; rtol=1f-5)
    @test isapprox((a + b) .* v, a .* v + b.* v; rtol=1f-5)
    @test isapprox(c2*v, c2*vec(v.data))
    @test isapprox(c2\v, c2\vec(v.data))
    @test isapprox(c + v, 1 .+ v)
    @test isapprox(c - v, 1 .- v)
    @test isapprox(v + c, 1 .+ v)
    @test isapprox(v - c, v .- 1)
    @test isapprox(v * v, v.^2)
    @test isapprox(v / v, 0f0.*v .+1)

    a = zeros(Float32, n...)
    a .= v
    @test isapprox(a, v.data)
    v = PhysicalParameter(ones(Float32, n...), d, o)
    v[v .> 0] .= -1
    @test all(v .== -1)

    # Indexing
    u = PhysicalParameter(randn(n), d, o)
    u[:] .= 0f0
    @test norm(u) == 0f0

    u[1] = 1f0
    @test norm(u) == 1f0
    @test u.data[1] == 1f0

    u[1:10] .= 1:10
    @test norm(u, 1) == 55
    @test u[1:10] == 1:10
    @test u.data[1:10] == 1:10
    @test u[11] == 0f0

    if nd == 2
        tmp = randn(u[1:10, :].n)
        u[1:10, :] .= tmp
        @test u.data[1:10, :] == tmp
        @test size(u[:, 1]) == (n[1],)
        @test size(u[1, :]) == (n[2],)
        @test u[2:end, 2:3].n == (n[1]-1, 2)
        u[:, 1] .= ones(Float32, n[1])
        @test all(u.data[:, 1] .== 1f0)
    else
        tmp = randn(u[1:10, :, :].n)
        u[1:10, :, :] .= tmp
        @test u.data[1:10, :, :] == tmp
        @test size(u[:, :, 1]) == (n[1], n[2])
        @test size(u[:, 1, :]) == (n[1], n[3])
        @test size(u[1, :, :]) == (n[2], n[3])
        @test size(u[1, 2:3, 1:end]) == (2, n[3])
        @test u[1:2, 2:3, 1:end].n == (2, 2, n[3])
        u[:, 1:end-2, 1] .= ones(Float32, n[1])
        @test all(u.data[:, 1:end-2, 1] .== 1f0)
    end

    # Test that works with JOLI
    A = joEye(size(v, 1))
    u2 = A*u
    @test isequal(u2, vec(u.data))
    u2 = [A;A]*u
    @test length(u2) == 2 * length(u)
    @test u2[1:length(u)] == vec(u.data)
    @test u2[length(u)+1:end] == vec(u.data)

end


text:
mutable struct POMCPOWPlanner{P,NBU,C,NA,SE,IN,IV,SolverType} <: Policy
    solver::SolverType
    problem::P
    node_sr_belief_updater::NBU
    criterion::C
    next_action::NA
    solved_estimate::SE
    init_N::IN
    init_V::IV
    tree::Union{Nothing, POMCPOWTree} # this is just so you can look at the tree later
end

function POMCPOWPlanner(solver, problem::POMDP)
    POMCPOWPlanner(solver,
                  problem,
                  solver.node_sr_belief_updater,
                  solver.criterion,
                  solver.next_action,
                  convert_estimator(solver.estimate_value, solver, problem),
                  solver.init_N,
                  solver.init_V,
                  nothing)
end

Random.seed!(p::POMCPOWPlanner, seed) = Random.seed!(p.solver.rng, seed)

function action_info(pomcp::POMCPOWPlanner{P,NBU}, b; tree_in_info=false) where {P,NBU}
    A = actiontype(P)
    info = Dict{Symbol, Any}()
    tree = make_tree(pomcp, b)
    pomcp.tree = tree
    local a::A
    try
        a = search(pomcp, tree, info)
        if pomcp.solver.tree_in_info || tree_in_info
            info[:tree] = tree
        end
    catch ex
        a = convert(A, default_action(pomcp.solver.default_action, pomcp.problem, b, ex))
    end
    return a, info
end

action(pomcp::POMCPOWPlanner, b) = first(action_info(pomcp, b))

function make_tree(p::POMCPOWPlanner{P, NBU}, b) where {P, NBU}
    S = statetype(P)
    A = actiontype(P)
    O = obstype(P)
    B = belief_type(NBU,P)
    return POMCPOWTree{B, A, O, typeof(b)}(b, 2*p.solver.tree_queries)
end

function search(pomcp::POMCPOWPlanner, tree::POMCPOWTree, info::Dict{Symbol,Any}=Dict{Symbol,Any}())
    all_terminal = true
    # gc_enable(false)
    i = 0
    start_us = CPUtime_us()
    for i in 1:pomcp.solver.tree_queries
        s = rand(pomcp.solver.rng, tree.root_belief)
        if !POMDPs.isterminal(pomcp.problem, s)
            max_depth = min(pomcp.solver.max_depth, ceil(Int, log(pomcp.solver.eps)/log(discount(pomcp.problem))))
            simulate(pomcp, POWTreeObsNode(tree, 1), s, max_depth)
            all_terminal = false
        end
        if CPUtime_us() - start_us >= pomcp.solver.max_time*1e6
            break
        end
    end
    info[:search_time_us] = CPUtime_us() - start_us
    info[:tree_queries] = i

    if all_terminal
        throw(AllSamplesTerminal(tree.root_belief))
    end

    best_node = select_best(pomcp.solver.final_criterion, POWTreeObsNode(tree,1), pomcp.solver.rng)

    return tree.a_labels[best_node]
end


text:
using QuasiGeostrophy, LinearAlgebra, Test, FFTW, BenchmarkTools, Plots

# boiler plate definitions
Ωxy = S¹(0,2π) × S¹(0,2π)
Nx = 2^6; Ny = 2^6;
fourier_grid = FourierGrid((Nx, Ny), Ωxy)
x, y = fourier_grid.grid
kx, ky = fourier_grid.wavenumbers
u = @. 1 * sin(x) * sin(y) + 0im
v = @. 1 * cos(x) * cos(y) + 0im
û = copy(u)
v̂ = copy(v)

FFTW.set_num_threads(Threads.nthreads())
P = plan_fft(u) # typeof(iP) <: AbstractFFTs.ScaledPlan
iP = plan_ifft(u)

uv = P * (u .* v)
norm(uv - convolve(P * u, P * v, P, iP)) / norm(uv)

## Define Parameter
const Re  = 100.0
const Δt = 1.0 / Re
const G  = 4.0

## Define forcing
forcing = @. G * sin(5 * x) * sin(5 * y) + 0im
f̂ = copy(forcing)
mul!(f̂, P, forcing)

## Define Operators
∂x = FourierOperator(im .* kx)
∂y = FourierOperator(im .* ky)
Δ = ∂x^2 + ∂y^2
Δ⁻¹ = inv(Δ)

function get_pressure(u_rhs, v_rhs, ∂x, ∂y, G)
    return G(∂x(u_rhs) + ∂y(v_rhs))
end

function closure_rhs(∂x, ∂y, Δ⁻¹, P, iP, forcing)
    function rhs(state)
        u = state[1]
        v = state[2]
        uu = convolve(u, u, P, iP)
        uv = convolve(u, v, P, iP)
        vv = convolve(v, v, P, iP)
        u_rhs = - ∂x(uu) - ∂y(uv) +  forcing
        v_rhs = - ∂x(uv) - ∂y(vv)
        p = get_pressure(u_rhs, v_rhs, ∂x, ∂y, Δ⁻¹)
        u_rhs -= ∂x(p)
        v_rhs -= ∂y(p)
        return [u_rhs, v_rhs]
    end
end

rhs = closure_rhs(∂x, ∂y, Δ⁻¹, P, iP, forcing)

function closure_lhs(Δ, Re, Δt)
    function lhs()
        u_lhs = 1 + -(Δt / Re) * Δ
        return [u_lhs, u_lhs]
    end
end

lhs = closure_lhs(Δ, Re, Δt)
operators = lhs()
lhs_operator = [inv(operator) for operator in operators]

function imex_step!(state, lhs_operator, rhs, Δt)
    rhs_state = rhs(state)
    current_state = copy(state)
    for i in eachindex(state)
        tmp = lhs_operator[i](current_state[i] .+ Δt .* rhs_state[i])
        @. state[i] =  tmp
    end
end

##
contourf(x[:], y[:], real.(forcing)', linewidth = 0, color = :thermometer)
##
u = @. 1 * sin(x) * sin(y) + 0im
v = @. 1 * cos(x) * cos(y) + 0im
û = copy(u)
v̂ = copy(v)
mul!(û, P, u)
mul!(v̂, P, v)
# û, v̂ = copy(good_start)
state = [û, v̂]
## Checking Incompressibility
rhs_tmp = rhs(state)
maximum(abs.(∂x(rhs_tmp[1]) + ∂y(rhs_tmp[2])))
##
for i in 1:20000
    imex_step!(state, lhs_operator, rhs, Δt)
    if (i%1000) == 0
        println(i)
        ω̂ = ∂x(v̂) - ∂y(û)
        mul!(u, iP, ω̂)
        local p1 = contourf(x[:],y[:], real.(u)', linewidth = 0, color = :thermometer)
        display(p1)
        println("checking conergence to steady state")
        println(norm(rhs(state) + (1/Re) * Δ.(state)) / (norm(rhs(state))))
        sleep(0.01)
    end
end
# good_start = copy(state)
## check steady state
# include(pwd() * "/sandbox/two_d_navier_stokes.jl")

max_stars_repo_name: upiterbarg/QuasiGeostrophy

max_issues_repo_name: sandreza/QuasiGeostrophy

max_forks_repo_name: upiterbarg/QuasiGeostrophy


text:
using Parameters, NLsolve, SpecialFunctions, Distributions, Roots, ArbNumerics, StatsBase, LsqFit, PoissonRandom, SparseArrays, Distributed, CSV, SharedArrays, JLD2, DataFrames, ProgressMeter

import GZip: open
import Parsers: parse
import OrderedCollections: OrderedDict
import FastaIO: readfasta

PATH = "/home/jmurga/.julia/dev/Analytical/src/"
include(PATH * "parameters.jl")
include(PATH * "fixations.jl")
include(PATH * "polymorphism.jl")
include(PATH * "summaryStatistics.jl")
include(PATH * "rates.jl")
include(PATH * "inferTools.jl")
include(PATH * "readFasta.jl")
include(PATH * "methods.jl")

param = parameters()
convolutedSamples = binomialDict()
binomOp!(param,convolutedSamples.bn)

function gettingRates(param::parameters,convolutedSamples::binomialDict)

	##############################################################
	# Accounting for positive alleles segregating due to linkage #
	##############################################################

	# Fixation
	fN       = param.B*fixNeut(param)
	fNeg     = param.B*fixNegB(param,0.5*param.pposH+0.5*param.pposL)
	fPosL    = fixPosSim(param,param.gL,0.5*param.pposL)
	fPosH    = fixPosSim(param,param.gH,0.5*param.pposH)

	ds       = fN
	dn       = fNeg + fPosL + fPosH

	## Polymorphism	## Polymorphism
	neut::Array{Float64,1} = DiscSFSNeutDown(param,convolutedSamples.bn[param.B])
	# neut = param.neut[param.B]

	selH::Array{Float64,1} = DiscSFSSelPosDown(param,param.gH,param.pposH,convolutedSamples.bn[param.B])
	selL::Array{Float64,1} = DiscSFSSelPosDown(param,param.gL,param.pposL,convolutedSamples.bn[param.B])
	selN::Array{Float64,1} = DiscSFSSelNegDown(param,param.pposH+param.pposL,convolutedSamples.bn[param.B])
	tmp = cumulativeSfs(hcat(neut,selH,selL,selN),false)
	splitColumns(matrix::Array{Float64,2}) = (view(matrix, :, i) for i in 1:size(matrix, 2));

	neut, selH, selL, selN = splitColumns(tmp)
	sel = (selH+selL)+selN

	## Outputs
	α = @. 1 - (ds/dn) * (sel/neut)

	##################################################################
	# Accounting for for neutral and deleterious alleles segregating #
	##################################################################
	## Fixation
	#=fN_nopos       = fN*(param.thetaMidNeutral/2.)*param.TE*param.NN
	fNeg_nopos     = fNeg*(param.thetaMidNeutral/2.)*param.TE*param.NN
	fPosL_nopos    = fPosL*(param.thetaMidNeutral/2.)*param.TE*param.NN
	fPosH_nopos    = fPosH*(param.thetaMidNeutral/2.)*param.TE*param.NN

	ds_nopos       = fN_nopos
	dn_nopos       = fNeg_nopos + fPosL_nopos + fPosH_nopos
	dnS_nopos      = dn_nopos - fPosL_nopos

	## Polymorphism
	sel_nopos = selN

	## Outputs
	αW         = param.alLow/param.alTot
	α_nopos    = @. 1 - (ds_nopos/dn_nopos) * (sel_nopos/neut)=#

	##########
	# Output #
	##########

	#=alphas = round.(vcat(α_nopos[param.dac[end]] * αW , α_nopos[param.dac[end]] * (1 - αW), α_nopos[param.dac[end]]), digits=5)=#
	analyticalValues::Array{Float64,2} = cat(param.B,(fPosL+fPosH),param.alLow,param.alTot,param.gamNeg,param.gL,param.gH,param.al,param.be,neut[param.dac],sel[param.dac],ds,dn,fPosL,fPosH,dims=1)'

	return (analyticalValues)
end

function iterRates(;param::parameters,convolutedSamples::binomialDict,alTot::Float64,alLow::Float64,gH::Int64,gL=Int64,gamNeg::Int64,afac::Float64)

	# Matrix and values to solve
	dm 			= 1
	param.al    = afac; param.be = abs(afac/gamNeg);
	param.alLow = alLow; param.alTot = alTot;
	param.gH    = gH;param.gL = gL; param.gamNeg = gamNeg

	# Solve probabilites without B effect to achieve α value
	param.B = 0.999
	setThetaF!(param)
	setPpos!(param)

	r = zeros(size(param.bRange,2) * dm,(size(param.dac,1) * 2) + 13)
	for j in eachindex(param.bRange)
		param.B = param.bRange[j]
		# Solve mutation given a new B value.
		setThetaF!(param)
		# Solven given same probabilites probabilites ≠ bgs mutation rate.
		#x,y,z::Array{Float64,2} = alphaByFrequencies(param,divergence,sfs,dac)
		@inbounds r[j,:] = gettingRates(param,convolutedSamples)
	end
	r[:,2] = @. r[:,2]/r[end,2]
	return r
end

function rates(;param::parameters,convolutedSamples::binomialDict,gH::Array{Int64,1},gL::Array{Int64,1},gamNeg::Array{Int64,1},shape::Float64=0.184,iterations::Int64,output::String)

	#=iterations=1;gamNeg=[-500];gL=[5];gH=[500];shape=0.184=#
	fac    = rand(-2:0.05:2,iterations,2)
	afac   = @. param.al*(2^fac[:,1])
	#=bfac   = @. (param.al/param.be)*(2^fac[:,2])=#

	lfac   = rand(0.05:0.05:0.9,iterations)
	nTot   = rand(0.1:0.01:0.9,iterations)

	nLow   = @. nTot * lfac
	nParam = [param for i in 1:iterations];
	nBinom = [convolutedSamples for i in 1:iterations];
	ngh    = rand(repeat(gH,iterations),iterations);
	ngl    = rand(repeat(gL,iterations),iterations);
	ngamNeg    = rand(repeat(gamNeg,iterations),iterations);

	# Estimations to thread pool
	out    = SharedArray{Float64,3}(size(param.bRange,2),(size(param.dac,1) *2) + 13,iterations)
	@sync @distributed for i in eachindex(afac)
		@inbounds out[:,:,i] = iterRates(param = nParam[i],convolutedSamples=nBinom[i],alTot = nTot[i], alLow = nLow[i],gH=ngh[i],gL=ngl[i],gamNeg=ngamNeg[i],afac=afac[i]);
	end

	df = vcat(eachslice(out,dims=3)...);
	models = DataFrame(df[:,1:9],[:B,:hri,:alLow,:alTot,:gamNeg,:gL,:gH,:al,:be])

    neut   = df[:,10:9+size(param.dac,1)]
    sel    = df[:,10+size(param.dac,1):9+size(param.dac,1)*2]
    dsdn   = df[:,end-3:end]

	JLD2.jldopen(output, "a+") do file
        file[string(param.N)* "/" * string(param.n) * "/models"] = models
        file[string(param.N)* "/" * string(param.n) * "/neut"]   = neut
        file[string(param.N)* "/" * string(param.n) * "/sel"]    = sel
        file[string(param.N)* "/" * string(param.n) * "/dsdn"]   = dsdn
        #=file[string(param.N)* "/" * string(param.n) * "/alphas"] = alphas=#
        file[string(param.N)* "/" * string(param.n) * "/dac"]    = param.dac
	end

	return df
end

function summaryStatsFromRates(;param::parameters,rates::JLD2.JLDFile,divergence::Array,sfs::Array,summstatSize::Int64,replicas::Int64)

	tmp    = rates[string(param.N) * "/" * string(param.n)]
	idx    = StatsBase.sample.(fill(1:size(tmp["neut"],1),replicas),fill(summstatSize,replicas),replace=false)

	models = Array.(map(x -> view(tmp["models"],x,:), idx))
	neut   = Array.(map(x -> view(tmp["neut"],x,:), idx))
	sel    = Array.(map(x -> view(tmp["sel"],x,:), idx))
	neut   = Array.(map(x -> view(tmp["neut"],x,:), idx))
	dsdn   = Array.(map(x -> view(tmp["dsdn"],x,:), idx))

	expectedValues =  progress_pmap(samplingFromRates,models,sfs,divergence,neut,sel,dsdn);

	return(expectedValues)
end

function samplingFromRates(m::Array,s::Array,d::Array,nt::Array,sl::Array,x::Array)
	ds      = x[:,1]
	dn      = x[:,2]
	dweak   = x[:,3]
	dstrong = x[:,4]
	gn      = abs.(m[:,5])
	sh      = round.(m[:,end-1],digits=5)
	hri     = m[:,2]

	alxSummStat, alphasDiv, expectedDn, expectedDs, expectedPn, expectedPs = sampledAlpha(d=d,afs=s,λdiv=[ds,dn,dweak,dstrong],λpol=[permutedims(nt),permutedims(sl)])
	expectedValues = hcat(round.(alphasDiv,digits=5),gn,sh,alxSummStat)
end

function sampledAlpha(;d::Array,afs::Array,λdiv::Array,λpol::Array)

	#=pn = λpol[:,2]
	ps = λpol[:,1]=#

	## Outputs
	alphas, expDn, expDs = poissonFixation(observedValues=d,λds=λdiv[1],λdn=λdiv[2],λweak=λdiv[3],λstrong=λdiv[4],hri=hri)
	expPn, expPs         = poissonPolymorphism(observedValues=afs,λps=λpol[1],λpn=λpol[2])

	## Alpha from expected values. Used as summary statistics
	αS = @. round(1 - ((expDs/expDn) * (expPn/expPs)'),digits=5)

	return αS,alphas,expDn,expDs,expPn,expPs
end

function poissonFixation(;observedValues::Array, λds::Array, λdn::Array,λweak::Array,λstrong::Array,hri::Array)

    ds        = @. λds / (λds + λdn)
    dn        = @. λdn / (λds + λdn)
    dweak     = @. λweak / (λds + λdn)
    dstrong   = @. λstrong / (λds + λdn)
    dpositive = @. (λweak + λstrong) / (λds + λdn)
    dnohri    = @. ((λweak + λstrong)/hri) / (λds + λdn)

	sampledDs     = PoissonRandom.pois_rand.(ds .* observedValues)
	sampledDn     = PoissonRandom.pois_rand.(dn .* observedValues)
	sampledWeak   = PoissonRandom.pois_rand.(dweak .* observedValues)
	sampledStrong = PoissonRandom.pois_rand.(dstrong .* observedValues)
	sampledPositive = PoissonRandom.pois_rand.(dpositive .* observedValues)
	sampledNoHri = PoissonRandom.pois_rand.(dnohri .* observedValues)

	h = sampledPositive./sampledNoHri

	h = 1 .- ifelse.(h .> 1,1,h)

	alphas = @. [sampledWeak/sampledDn sampledStrong/sampledDn sampledPositive/sampledDn h]

	out = alphas,sampledDn, sampledDs
	return out
end


text:
function find_optimal_iteration(gm::GlobalModel)
    if any(bbl.equality == true for bbl in gm.bbls)
        @warn "GlobalModel $(gm.name) has equalities. These aren't guaranteed to be feasible. "
    end
    noteqs = [bbl.equality == false for bbl in gm.bbls]
    bbcs = [bbl isa BlackBoxClassifier for bbl in gm.bbls]
    bbrs = [bbl for bbl in gm.bbls if bbl isa BlackBoxRegressor]
    if isempty(bbrs)
        for i = sortperm(gm.cost)
            feas = gm.feas_history[i] .* noteqs .* bbcs
            if !any(feas .<= -1e-10)
                @info "Iteration $(i) was optimal."
                return i
            end
        end
    else
        for i = sortperm(bbrs[1].actuals)
            feas = gm.feas_history[i] .* noteqs .* bbcs
            if !any(feas .<= -1e-10)
                @info "Iteration $(i) was optimal."
                return i
            end
        end
    end
    @info "None of the solutions fit the feasibility criteria. "
    return
end

# Iterative training algorithm
    # if !isempty(bbrs)
    #     @assert length(bbrs) == 1
    #     bbr = bbrs[1]
    #     learn_constraint!(bbr, "upper" => minimum(bbr.actuals))
    #     update_tree_constraints!(gm, bbr)
    #     upper_bound_sample(bbr)
    #     learn_constraint!(bbr, "reg" => minimum(bbr.actuals), regression_sparsity = 0, max_depth = 5)
    #     update_tree_constraints!(gm, bbr)
    #     update_leaf_vexity(bbr)
    #     optimize!(gm)
    #     while abs(gm.cost[end] - gm.cost[end-1]) > get_param(gm, :abstol)
    #         add_infeasibility_cuts!(gm)
    #         optimize!(gm)
    #     end
    #     while minimum(bbr.actuals) < bbr.thresholds[active_lower_tree(bbr)].second
    #         for bbc in bbcs # to clear already present cuts
    #             clear_tree_constraints!(gm, bbc)
    #             add_tree_constraints!(gm, bbc)
    #         end
    #         learn_constraint!(bbr, "upper" => minimum(bbr.actuals))
    #         update_tree_constraints!(gm, bbr)
    #         upper_bound_sample(bbr)
    #         learn_constraint!(bbr, "reg" => minimum(bbr.actuals), regression_sparsity = 0, max_depth = 5)
    #         update_tree_constraints!(gm, bbr)
    #         update_leaf_vexity(bbr)
    #         try
    #             optimize!(gm)
    #         catch
    #             break
    #         end
    #         while abs(gm.cost[end] - gm.cost[end-1]) > get_param(gm, :abstol)
    #             add_infeasibility_cuts!(gm)
    #             try
    #                 optimize!(gm)
    #             catch
    #                 @info "Optimization terminated due to infeasibility."
    #                 break
    #             end
    #         end
    #     end
    #     @info "Optimization terminated."
    #     it = find_optimal_iteration(gm)
    #     @info "Elapsed time: $(round(time()-t1, sigdigits=5)) seconds."
    #     @info "Optimal cost: $(round(bbr.actuals[it], sigdigits = 7))"
    # else


text:
###############################################
#                                             #
# Helper functions:                           #
# Not of immediate interest, but useful       #
# for solving the model                       #
#                                             #
###############################################

# Included:
# - Bellman
# - Bellman_NLopt
# - CC
# - VTerm
# - YY_noep
# - YY

##########################################################################################
# Computes (negative of) Bellman equation
function Bellman(Par,Grid,S,I,R,V0,x)
    q = x[1];
    lam = x[2];
    mu = x[3];
    IN = max(1-S-R,0);
    if Par.TEST == 1
        Del = q.*Par.theti;
    else
        Del = 0;
    end
    if Par.ISO == 1
        LI1 = lam.*((1-q.*Par.thets).*S+(1-q.*Par.theti).*I)+mu.*(1-Par.delta).*R;
        LI2 = (1-lam).*((1-q.*Par.thets).*S+(1-q.*Par.theti).*I)+(1-mu).*(1-Par.delta).*R;
        if LI1>0
            Ipr1 = (lam.^2 .* Par.pi1 .* (1-q.*Par.theti).*I)./LI1;
        else
            Ipr1 = 0;
        end
        if LI2>0
            Ipr2 = ((1-lam).^2 .* Par.pi2 .* (1-q.*Par.theti).*I)./LI2;
        else
            Ipr2 = 0;
        end
        Ipr = Ipr1 + Ipr2;
        Ipr = min(Ipr,1);
        Ip = min(max((1-Par.tautilde).*(1-Par.gamma).*I+(1-q.*Par.thets).*S.*Ipr,0),1);
        Sp = min(max((1-(1-q.*Par.thets).*Ipr).*S,0),1);
        Rp = max(min(R + Par.gamma.*IN,1),0);
    else
        LI1 = lam.*(S+(1-Del).*I)+mu.*(1-Par.delta).*R;
        LI2 = (1-lam).*(S+(1-Del).*I)+(1-mu).*(1-Par.delta).*R;
        if LI1>0
            Ipr1 = (lam.^2 .* Par.pi1 .* (1-Del).*I)./LI1;
        else
            Ipr1 = 0;
        end
        if LI2>0
            Ipr2 = ((1-lam).^2 .* Par.pi2 .* (1-Del).*I)./LI2;
        else
            Ipr2 = 0;
        end
        Ipr = Ipr1 + Ipr2;
        Ipr = min(Ipr,1);
        Ip = min(max((1-Del).*(1-Par.tautilde).*(1-Par.gamma).*I+S.*Ipr,0),1);
        Sp = min(max((1-Ipr).*S,0),1);
        Rp = max(min(R + Par.gamma.*IN,1),0);
    end

    Y = YY(Par,S,I,R,q,lam,mu);
    C1,C2 = CC(Par,S,I,R,q,lam,mu);
    OO = max(Y-C1,0);

    if OO>0
        #grid = SimplexGrid(Grid.S, Grid.I, Grid.R);
        grid = RectangleGrid(Grid.S, Grid.I, Grid.R);
        gridData = V0;
        V_interpolated = ((1-Par.beta).*(log(OO) - C2) + Par.beta .* interpolate(grid, gridData, [Sp, Ip, Rp]));
        if isnan(V_interpolated)
           V = -VTerm(Par,0,0,0);
        else
            V = -min(VTerm(Par,0,0,0), V_interpolated);
        end
    else
        V = -(-100000000);
    end

    return V
end

##########################################################################################
# Wrapper for Bellman, as expected by NLopt optimizer
function Bellman_NLopt(Par, Grid, S, I, R, V00, x::Vector, g::Vector)
    gradient = grad(backward_fdm(2, 1), z -> Bellman(Par, Grid, S, I, R, V00, z), x)[1]
    if length(g) > 0
        g[1] = gradient[1]
        g[2] = gradient[2]
        g[3] = gradient[3]
    end
    return Bellman(Par, Grid, S, I, R, V00, x)
end

##########################################################################################
# Computes total cost (infection+testing+death)
function CC(Par,S,I,R,q,lam,mu)
    IN = max(1-S-R,0);
    if Par.TEST == 1
        x1 = Par.c.*(q.*min((Par.thets.*S+Par.theti.*I),1)).^(1+Par.nu)./(1+Par.nu);
        x2 = Par.z.*(IN-I)+Par.d.*Par.gamma.*Par.delta.*IN;
    else
        x1 = 0;
        x2 = Par.z.*(IN-I)+Par.d.*Par.gamma.*Par.delta.*IN;
    end
    return [x1, x2]
end

##########################################################################################
# Computes terminal value function
function VTerm(Par,S,I,R)
    lam = ((Par.B1./Par.B2).^Par.sigma)./(1+(Par.B1./Par.B2).^Par.sigma);
    y = (Par.B1.*lam.^((Par.sigma-1)./(Par.sigma))+Par.B2.*(1-lam).^((Par.sigma-1)./(Par.sigma))).^(Par.sigma./(Par.sigma-1));
    return log(y.*(1-Par.delta.*R));
    end

##########################################################################################
# Computes output in the no epidemic steady state:
# Using Par:
function YY_noep(Par, L)
    lam = ((Par.B1./Par.B2).^Par.sigma)./(1+(Par.B1./Par.B2).^Par.sigma);
    y = L.*(Par.B1.*lam.^((Par.sigma-1)./(Par.sigma))+Par.B2.*(1-lam).^((Par.sigma-1)./(Par.sigma))).^(Par.sigma./(Par.sigma-1));
    return y
end
# Using parameters of interest directly
function YY_noep(B1, B2, sigma, L)
    lam = ((B1./B2).^sigma) ./ (1+(B1./B2).^sigma);
    y = L.*(B1.*lam.^((sigma-1)./(sigma))+B2.*(1-lam).^((sigma-1)./(sigma))).^(sigma./(sigma-1));
    return y
end

##########################################################################################
# Computes output
function YY(Par,S,I,R,q,lam,mu)
    if Par.ISO == 0 && Par.TEST == 0
        L1 = (lam.*(S+Par.xi.*I)+mu.*(1-Par.delta).*R);
        L2 = ((1-lam).*(S+Par.xi.*I)+(1-mu).*(1-Par.delta).*R);
        y = (Par.B1.*L1.^((Par.sigma-1)/Par.sigma)+Par.B2.*L2.^((Par.sigma-1)/Par.sigma)).^(Par.sigma./((Par.sigma-1)));
    elseif Par.ISO == 1
        L1 = (lam.*(S.*(1-q.*Par.thets)+Par.xi.*(1-q.*Par.theti).*I)+mu.*(1-Par.delta).*R);
        L2 = ((1-lam).*(S.*(1-q.*Par.thets)+Par.xi.*(1-q.*Par.theti).*I)+(1-mu).*(1-Par.delta).*R);
        y = (Par.B1.*L1.^((Par.sigma-1)/Par.sigma)+Par.B2.*L2.^((Par.sigma-1)/Par.sigma)).^(Par.sigma./((Par.sigma-1)));
    else
        L1 = (lam.*(S+Par.xi.*(1-q.*Par.theti).*I)+mu.*(1-Par.delta).*R);
        L2 = ((1-lam).*(S+Par.xi.*(1-q.*Par.theti).*I)+(1-mu).*(1-Par.delta).*R);
        y = (Par.B1.*L1.^((Par.sigma-1)/Par.sigma)+Par.B2.*L2.^((Par.sigma-1)/Par.sigma)).^(Par.sigma./((Par.sigma-1)));
    end
    return y
end


text:
abstract type IntegralOperator <: Operator end

"""
    blockassembler(operator, test_space, trial_space) -> assembler

Return a callable object for the creation of blocks within a BEM matrix.

This function performs all tasks common to the assembly of several blocks within
a single boundary element matrix. The return value can be used to generate blocks
by calling it as follows:

    assembler(I,J,storefn)

where `I` and `J` are arrays of indices in `test_space` and `trial_space`, respectively,
corresponding to the rows and columns of the desired block.

Note that the block will be constructed in compressed form, i.e. the rows and columns
of the store that are written into are the positions within `I` and `J` (as opposed
to the positions within `1:numfunctions(test_space)` and `1:numfunctions(trial_space)`).
In particular the size of the constructed block will be `(length(I), length(J))`.

This last property allows the assembly of permutations of the BEM matrix by supplying
for `I` and `J` permutations of `1:numfunctions(test_space)` and
`1:numfunctions(trial_space)`.
"""
function blockassembler end

"""
    quadrule(operator,test_refspace,trial_refspace,p,test_element,q_trial_element, qd)

Returns an object that contains all the dynamic (runtime) information that
defines the integration strategy that will be used by `momintegrals!` to compute
the interactions between the local test/trial functions defined on the specified
geometric elements. The indices `p` and `q` refer to the position of the test
and trial elements as encountered during iteration over the output of
`geometry`.

The last argument `qd` provides access to all precomputed data required for
quadrature. For example it might be desirable to precompute all the quadrature
points for all possible numerical quadrature schemes that can potentially be
required during matrix assembly. This makes sense, since the number of point is
order N (where N is the number of faces) but these points will appear in N^2
computations. Precomputation requires some extra memory but can save a lot on
computation time.
"""
function quadrule end

"""
  elements(geo)

Create an iterable collection of the elements stored in `geo`. The order in which
this collection produces the elements determines the index used for lookup in the
data structures returned by `assemblydata` and `quaddata`.
"""
elements(geo) = [chart(geo,cl) for cl in cells(geo)]

elements(sp::Space) = elements(geometry(sp))

"""
    assemblechunk!(biop::IntegralOperator, tfs, bfs, store)

Computes the matrix of operator biop wrt the finite element spaces tfs and bfs
"""
function assemblechunk!(biop::IntegralOperator, tfs::Space, bfs::Space, store;
        quadstrat=defaultquadstrat(biop, tfs, bfs))

    test_elements, tad = assemblydata(tfs)
    bsis_elements, bad = assemblydata(bfs)

    tshapes = refspace(tfs); num_tshapes = numfunctions(tshapes)
    bshapes = refspace(bfs); num_bshapes = numfunctions(bshapes)

    qd = quaddata(biop, tshapes, bshapes, test_elements, bsis_elements, quadstrat)
    zlocal = zeros(scalartype(biop, tfs, bfs), 2num_tshapes, 2num_bshapes)

    if !CompScienceMeshes.refines(tfs.geo, bfs.geo)
        assemblechunk_body!(biop,
            tshapes, test_elements, tad,
            bshapes, bsis_elements, bad,
            qd, zlocal, store; quadstrat)
    else
        @info "assemblechunk for nested meshes"
        assemblechunk_body_nested_meshes!(biop,
            tshapes, test_elements, tad,
            bshapes, bsis_elements, bad,
            qd, zlocal, store; quadstrat)
    end
end

function assemblechunk_body!(biop,
        test_shapes, test_elements, test_assembly_data,
        trial_shapes, trial_elements, trial_assembly_data,
        qd, zlocal, store; quadstrat)

    myid = Threads.threadid()
    myid == 1 && print("dots out of 10: ")
    todo, done, pctg = length(test_elements), 0, 0
    for (p,tcell) in enumerate(test_elements)
        for (q,bcell) in enumerate(trial_elements)

        fill!(zlocal, 0)
        qrule = quadrule(biop, test_shapes, trial_shapes, p, tcell, q, bcell, qd, quadstrat)
        momintegrals!(biop, test_shapes, trial_shapes, tcell, bcell, zlocal, qrule)
        I = length(test_assembly_data[p])
        J = length(trial_assembly_data[q])
        for j in 1 : J, i in 1 : I
            zij = zlocal[i,j]
            for (n,b) in trial_assembly_data[q][j]
                zb = zij*b
                for (m,a) in test_assembly_data[p][i]
                    store(a*zb, m, n)
        end end end end

        done += 1
        new_pctg = round(Int, done / todo * 100)
        if new_pctg > pctg + 9
            myid == 1 && print(".")
            pctg = new_pctg
    end end
    myid == 1 && println("")
end

function assemblechunk_body_nested_meshes!(biop,
        test_shapes, test_elements, test_assembly_data,
        trial_shapes, trial_elements, trial_assembly_data,
        qd, zlocal, store; quadstrat)

    myid = Threads.threadid()
    myid == 1 && print("dots out of 10: ")
    todo, done, pctg = length(test_elements), 0, 0
    for (p,tcell) in enumerate(test_elements)
        for (q,bcell) in enumerate(trial_elements)

            fill!(zlocal, 0)
            qrule = quadrule(biop, test_shapes, trial_shapes, p, tcell, q, bcell, qd, quadstrat)
            momintegrals_nested!(biop, test_shapes, trial_shapes, tcell, bcell, zlocal, qrule)
            I = length(test_assembly_data[p])
            J = length(trial_assembly_data[q])
            for j in 1 : J, i in 1 : I
                zij = zlocal[i,j]
                for (n,b) in trial_assembly_data[q][j]
                    zb = zij*b
                    for (m,a) in test_assembly_data[p][i]
                    store(a*zb, m, n)
        end end end end

        done += 1
        new_pctg = round(Int, done / todo * 100)
        if new_pctg > pctg + 9
            myid == 1 && print(".")
            pctg = new_pctg
        end end
    myid == 1 && println("")
end

function blockassembler(biop::IntegralOperator, tfs::Space, bfs::Space;
        quadstrat=defaultquadstrat(biop, tfs, bfs))

    test_elements, test_assembly_data,
        trial_elements, trial_assembly_data,
        quadrature_data, zlocals = assembleblock_primer(biop, tfs, bfs; quadstrat)

    if !CompScienceMeshes.refines(tfs.geo, bfs.geo)
        return (test_ids, trial_ids, store) -> begin
            assembleblock_body!(biop,
                tfs, test_ids,   test_elements,  test_assembly_data,
                bfs, trial_ids, trial_elements, trial_assembly_data,
                quadrature_data, zlocals, store; quadstrat)
        end
    else
        return (test_ids, trial_ids, store) -> begin
            assembleblock_body_nested!(biop,
                tfs, test_ids,   test_elements,  test_assembly_data,
                bfs, trial_ids, trial_elements, trial_assembly_data,
                quadrature_data, zlocals, store; quadstrat)
        end
    end
end

function assembleblock(operator::AbstractOperator, test_functions, trial_functions;
        quadstrat=defaultquadstrat(operator, test_functions, trial_functions))

    Z, store = allocatestorage(operator, test_functions, trial_functions)
    assembleblock!(operator, test_functions, trial_functions, store; quadstrat)

    sdata(Z)
end

function assembleblock!(biop::IntegralOperator, tfs::Space, bfs::Space, store;
        quadstrat=defaultquadstrat(biop, tfs, bfs))

    test_elements, tad, trial_elements, bad, quadrature_data, zlocals =
        assembleblock_primer(biop, tfs, bfs; quadstrat)

    active_test_dofs  = collect(1:numfunctions(tfs))
    active_trial_dofs = collect(1:numfunctions(bfs))

    assembleblock_body!(biop,
        tfs, active_test_dofs, test_elements, tad,
        bfs, active_trial_dofs, trial_elements, bad,
        quadrature_data, zlocals, store; quadstrat)
end

function assembleblock_primer(biop, tfs, bfs;
        quadstrat=defaultquadstrat(biop, tfs, bfs))

    test_elements, tad = assemblydata(tfs)
    bsis_elements, bad = assemblydata(bfs)

    tshapes = refspace(tfs); num_tshapes = numfunctions(tshapes)
    bshapes = refspace(bfs); num_bshapes = numfunctions(bshapes)

    qd = quaddata(biop, tshapes, bshapes, test_elements, bsis_elements, quadstrat)

    zlocals = Matrix{scalartype(biop, tfs, bfs)}[]

    for i in 1:Threads.nthreads()
        push!(zlocals, zeros(scalartype(biop, tfs, bfs), num_tshapes, num_bshapes))
    end

    return test_elements, tad, bsis_elements, bad, qd, zlocals
end

function assembleblock_body!(biop::IntegralOperator,
        tfs, test_ids, test_elements, test_assembly_data,
        bfs, trial_ids, bsis_elements, trial_assembly_data,
        quadrature_data, zlocals, store; quadstrat)

    test_shapes  = refspace(tfs)
    trial_shapes = refspace(bfs)

    # Enumerate all the active test elements
    active_test_el_ids  = Vector{Int}()
    active_trial_el_ids = Vector{Int}()

    test_id_in_blk  = Dict{Int,Int}()
    trial_id_in_blk = Dict{Int,Int}()

    for (i,m) in enumerate(test_ids);   test_id_in_blk[m] = i; end
    for (i,m) in enumerate(trial_ids); trial_id_in_blk[m] = i; end

    for m in test_ids,  sh in tfs.fns[m]; push!(active_test_el_ids,  sh.cellid); end
    for m in trial_ids, sh in bfs.fns[m]; push!(active_trial_el_ids, sh.cellid); end

    active_test_el_ids = unique!(sort!(active_test_el_ids))
    active_trial_el_ids = unique!(sort!(active_trial_el_ids))

    for p in active_test_el_ids
        tcell = test_elements[p]
        for q in active_trial_el_ids
            bcell = bsis_elements[q]

            fill!(zlocals[Threads.threadid()], 0)
            qrule = quadrule(biop, test_shapes, trial_shapes, p, tcell, q, bcell, quadrature_data, quadstrat)
            momintegrals!(biop, test_shapes, trial_shapes, tcell, bcell, zlocals[Threads.threadid()], qrule)

            for j in 1 : size(zlocals[Threads.threadid()],2)
                for i in 1 : size(zlocals[Threads.threadid()],1)
                    for (n,b) in trial_assembly_data[q,j]
                        n′ = get(trial_id_in_blk, n, 0)
                        n′ == 0 && continue
                        for (m,a) in test_assembly_data[p,i]
                            m′ = get(test_id_in_blk, m, 0)
                            m′ == 0 && continue
                            store(a*zlocals[Threads.threadid()][i,j]*b, m′, n′)
end end end end end end end

function assembleblock_body_nested!(biop::IntegralOperator,
    tfs, test_ids, test_elements, test_assembly_data,
    bfs, trial_ids, bsis_elements, trial_assembly_data,
    quadrature_data, zlocals, store; quadstrat)

    test_shapes  = refspace(tfs)
    trial_shapes = refspace(bfs)

    # Enumerate all the active test elements
    active_test_el_ids  = Vector{Int}()
    active_trial_el_ids = Vector{Int}()

    test_id_in_blk  = Dict{Int,Int}()
    trial_id_in_blk = Dict{Int,Int}()

    for (i,m) in enumerate(test_ids);   test_id_in_blk[m] = i; end
    for (i,m) in enumerate(trial_ids); trial_id_in_blk[m] = i; end

    for m in test_ids,  sh in tfs.fns[m]; push!(active_test_el_ids,  sh.cellid); end
    for m in trial_ids, sh in bfs.fns[m]; push!(active_trial_el_ids, sh.cellid); end

    active_test_el_ids = unique(sort(active_test_el_ids))
    active_trial_el_ids = unique(sort(active_trial_el_ids))

    for p in active_test_el_ids
        tcell = test_elements[p]
        for q in active_trial_el_ids
            bcell = bsis_elements[q]

            fill!(zlocals[Threads.threadid()], 0)
            qrule = quadrule(biop, test_shapes, trial_shapes, p, tcell, q, bcell, quadrature_data, quadstrat)
            momintegrals_nested!(biop, test_shapes, trial_shapes, tcell, bcell, zlocals[Threads.threadid()], qrule)

            for j in 1 : size(zlocals[Threads.threadid()],2)
                for i in 1 : size(zlocals[Threads.threadid()],1)
                    for (n,b) in trial_assembly_data[q,j]
                        n′ = get(trial_id_in_blk, n, 0)
                        n′ == 0 && continue
                        for (m,a) in test_assembly_data[p,i]
                            m′ = get(test_id_in_blk, m, 0)
                            m′ == 0 && continue
                            store(a*zlocals[Threads.threadid()][i,j]*b, m′, n′)
end end end end end end end

function assemblerow!(biop::IntegralOperator, test_functions::Space, trial_functions::Space, store;
        quadstrat=defaultquadstrat(biop, test_functions, trial_functions))

    test_elements = elements(geometry(test_functions))
    trial_elements, trial_assembly_data = assemblydata(trial_functions)

    test_shapes  = refspace(test_functions)
    trial_shapes = refspace(trial_functions)

    num_test_shapes  = numfunctions(test_shapes)
    num_trial_shapes = numfunctions(trial_shapes)

    quadrature_data = quaddata(biop, test_shapes, trial_shapes, test_elements, trial_elements,
        quadstrat)
    zlocal = zeros(scalartype(biop, test_functions, trial_functions),
        num_test_shapes, num_trial_shapes)

    @assert length(trial_elements) == numcells(geometry(trial_functions))
    @assert numfunctions(test_functions) == 1

    assemblerow_body!(biop,
        test_functions, test_elements, test_shapes,
        trial_assembly_data, trial_elements, trial_shapes,
        zlocal, quadrature_data, store; quadstrat)
end

function assemblerow_body!(biop,
    test_functions, test_elements, test_shapes,
    trial_assembly_data, trial_elements, trial_shapes,
    zlocal, quadrature_data, store; quadstrat)

    test_function = test_functions.fns[1]
    for shape in test_function
        p = shape.cellid
        i = shape.refid
        a = shape.coeff
        tcell = test_elements[p]
        for (q,bcell) in enumerate(trial_elements)

            fill!(zlocal, 0)
            qrule = quadrule(biop, test_shapes, trial_shapes, p, tcell, q, bcell, quadrature_data, quadstrat)
            momintegrals!(biop, test_shapes, trial_shapes, tcell, bcell, zlocal, qrule)

            for j in 1:size(zlocal,2)
                for (n,b) in trial_assembly_data[q,j]
                    store(a*zlocal[i,j]*b, 1, n)
end end end end end

function assemblecol!(biop::IntegralOperator, test_functions::Space, trial_functions::Space, store;
        quadstrat=defaultquadstrat(biop, test_functions, trial_functions))

    test_elements, test_assembly_data = assemblydata(test_functions)
    trial_elements = elements(geometry(trial_functions))

    test_shapes  = refspace(test_functions)
    trial_shapes = refspace(trial_functions)

    num_test_shapes  = numfunctions(test_shapes)
    num_trial_shapes = numfunctions(trial_shapes)

    quadrature_data = quaddata(biop, test_shapes, trial_shapes, test_elements, trial_elements, quadstrat)
    zlocal = zeros(
        scalartype(biop, test_functions, trial_functions),
        num_test_shapes, num_trial_shapes)

    @assert length(test_elements) == numcells(geometry(test_functions))
    @assert numfunctions(trial_functions) == 1

    assemblecol_body!(biop,
        test_assembly_data, test_elements,  test_shapes,
        trial_functions,   trial_elements, trial_shapes,
        zlocal, quadrature_data, store; quadstrat)
end

function assemblecol_body!(biop,
    test_assembly_data, test_elements, test_shapes,
    trial_functions, trial_elements, trial_shapes,
    zlocal, quadrature_data, store; quadstrat)

    trial_function = trial_functions.fns[1]
    for shape in trial_function
        q = shape.cellid
        j = shape.refid
        b = shape.coeff

        bcell = trial_elements[q]
        for (p,tcell) in enumerate(test_elements)

            fill!(zlocal, 0)
            qrule = quadrule(biop, test_shapes, trial_shapes, p, tcell, q, bcell, quadrature_data, quadstrat)
            momintegrals!(biop, test_shapes, trial_shapes, tcell, bcell, zlocal, qrule)

            for i in 1:size(zlocal,1)
                for (m,a) in test_assembly_data[p,i]
                    store(a*zlocal[i,j]*b, m, 1)
end end end end end


text:
@everywhere ROOT = pwd()
@everywhere include(joinpath(ROOT,"src_PowSysMod", "PowSysMod_body.jl"))
@everywhere include("para_fct.jl")

# NOTE: Launch in console using 'julia -p nbparathreads compute_GOC_all_scenario.jl'
# where the nb of cores is likely a good value for nbparathreads

folder = "Phase_0_IEEE14"
folder = "Phase_0_RTS96"
folder_path = joinpath("instances", "GOC", folder)

scenarios = sort(filter(x->!ismatch(r"\.", x), readdir(folder_path)), by=x->parse(split(x, "_")[2]))

date = Dates.format(now(), "yy_u_dd_HH_MM_SS")
save_folder = joinpath("knitro_runs", "para_run_$date")
mkdir(save_folder)

println("----------> Start para jobs")

# r = build_and_solve_GOC(folder, scenarios[1])
r = pmap(build_and_solve_GOC, [folder for i=1:length(scenarios)], scenarios, [save_folder for i=1:length(scenarios)])

println("----------> para jobs done")

println("length scenarios : $(length(r))")

# r = Dict(r)
#
# tasks_th = Dict()
# for (key, val) in r
#     if !haskey(tasks_th, val["thread_id"])
#         tasks_th[val["thread_id"]] = 0
#     end
#     tasks_th[val["thread_id"]] += 1
# end
# println("id count:")
# for (k, v) in tasks_th
#     println("$k → $v")
# end

max_stars_repo_name: Gimaju/GOC

max_issues_repo_name: Gimaju/GOC

max_forks_repo_name: Gimaju/GOC


text:
struct VLFResult
    fit::Any
    data::VirusLoadData
    names::Vector
    p0::Vector
end

heaviside(x::AbstractFloat) = ifelse(x < 0, zero(x), ifelse(x > 0, one(x), oftype(x,0.5)))

v₁(t, a₁, a₂, logVmax) = 1. + (10^logVmax - 1.0)*(tanh(6.0*(t - (a₁ + a₂)/2)/(a₂ - a₁)) - tanh(-3.0*(a₂ + a₁)/(a₂ - a₁)))/2

v₂(t, a₂, α) = 1. - heaviside(t - a₂) + heaviside(t - a₂)*exp(-α*(t - a₂))

v₃(t, b₁, b₂, logVmin) = 1. - (1.0 - 10^logVmin)*(tanh(6.0*(t - (b₁ + b₂)/2)/(b₂ - b₁)) - tanh(-3.0*(b₂ + b₁)/(b₂ - b₁)))/2

function assert_params(a₁, a₂, b₁, b₂, α, logVmax)
    if a₁ > a₂ @error("a₁ > a₂") end
    if a₂ > b₁ @error("a₂ > b₁") end
    if b₁ > b₂ @error("b₁ > b₂") end
end

function LogVirusLoadFunction(t, p, data::VirusLoadData)
    logVmin = -7.0
    a₁, a₂, b₁, b₂, α, logVmax = p
    log10.(v₁.(t, a₁, a₂, logVmax).*v₂.(t, a₂, α).*v₃.(t, b₁, b₂, logVmin))
end

function LogEmpiricalVirusLoadFunction(t, p, data::VirusLoadData)
    θ = minimum(data.v)
    logV = LogVirusLoadFunction(t, p, data::VirusLoadData)
    max.(logV, θ)
end

function get_bounds(data, p0)
    a₁, a₂, b₁, b₂, α, logVmax = p0
    a₁a₂ = (a₁ + a₂)/2
    a₂b₁ = (a₂ + b₁)/2
    b₁b₂ = (b₁ + b₂)/2
    tmin = 0.0
    tmax = maximum(data.t)
    vmax = maximum(data.v)
    lb   = [tmin, a₁a₂, a₂b₁, b₁b₂, 1e-8, vmax-3.0]
    ub   = [a₁a₂, a₂b₁, b₁b₂, tmax, 1e+4, vmax+3.0]
    return lb, ub
end

function fitVLF(data::VirusLoadData, p0::Vector)
    model(t, p) = LogEmpiricalVirusLoadFunction(t, p, data)
    lb, ub = get_bounds(data, p0)
    fit = curve_fit(model, data.t, data.v, p0, lower=lb, upper=ub)
    names = ["a₁", "a₂", "b₁", "b₂", "α", "logVmax"]
    VLFResult(fit, data, names, p0)
end

@recipe function f(result::VLFResult; empirical=false)
    tmin, tmax = 0.0, maximum(result.data.t)
    vmin, vmax = extrema(result.data.v)
    tt = Vector(range(tmin, tmax, step=1e-2))
    x := tt
    if empirical
        y := LogEmpiricalVirusLoadFunction(tt, result.fit.param, result.data)
    else
        y := LogVirusLoadFunction(tt, result.fit.param, result.data)
    end
    linewidth --> 4
    xaxis --> ("Time (days)", (tmin, tmax))
    yaxis --> (L"\log\,V(t)", (vmin-1, vmax+1))
    grid --> :none
    label --> "Virus load function"
    ()
end

function Base.summary(result::VLFResult)
    println(@sprintf "RSS = %.5e (convergence: %s)" sum(abs2, result.fit.resid) result.fit.converged)
    CI = try
        confidence_interval(result.fit, 0.05)
    catch
        false
    end
    for i in 1:length(result.fit.param)
        name = result.names[i]
        val = result.fit.param[i]
        p0 = result.p0[i]
        if CI == false
            println(@sprintf "  %s = %.3e (initial=%.3e)" name val p0)

        else
            CIl = CI[i][1]
            CIr = CI[i][2]
            println(@sprintf "  %s = %.3e (CI=(%.3e, %.3e), initial=%.3e)" name val CIl CIr p0)
        end
    end
end


text:
export make_LN_basis,
       make_LN1N2_basis,
       make_spin1_LSz_basis,
       make_spin1_L_basis,
       get_num_spin1_states

"""
    make_LN_basis(L::Int, N::Int)

Compute Fock basis states with L sites and N particles.
"""
function make_LN_basis(L::Int, N::Int)
    states = Vector{Int}(undef, binomial(L, N))
    tmp = 1
    for i=0:(1<<L)-1
        if count_ones(i) == N
            states[tmp] = i
            tmp += 1
        end
    end
    return states
end

"""
    make_LN1N2_basis(L::Int, N1::Int, N2::Int)

Compute Fock basis states with L sites, N1 particles of type 1 and N2 particles
of type 2 (e.g. spin up and spin down).
"""
function make_LN1N2_basis(L::Int, N1::Int, N2::Int)
    states = Vector{Int}(undef, binomial(L, N1)*binomial(L, N2))
    tmp = 1
    mask1 = 1<<L - 1
    mask2 = 1<<(2L) - 1 - mask1
    for i=0:(1<<(2L))-1
        if count_ones(i&mask1) == N1 && count_ones(i&mask2) == N2
            states[tmp] = i
            tmp += 1
        end
    end
    return states
end

"""
    make_spin1_LSz_basis(L::Int, Sz::Int)

Compute states with L spins s=1 and Sz magnetization.

There can be 3 states per individual site/spin. In the usual notation
these are: |1, 0>, |1, 1>, |1, -1>. Because we can't represent these
states with only one bit of information we map each site to the three
possible states:
* |1, 0> = (00)
* |1, 1> = (01)
* |1,-1> = (10)
For example, a state with three spins, with magnetization -1, 1, 0 in
each site, respectively, is represented as: 100100.
"""
function make_spin1_LSz_basis(L::Int, Sz::Int)
    # Number of states.
    num_states = 0
    l = iseven(L+Sz) ? L : L-1
    # Number of spins with m=1 and m=-1 in a chain of l sites.
    Nu = (l+Sz)>>1
    Nd = (l-Sz)>>1
    while l >= 0 && Nu >= 0 && Nd >= 0
        num_states += binomial(l, Nu)*binomial(L, l)
        l -= 2
        Nu -= 1
        Nd -= 1
    end

    states = zeros(Int, num_states)
    cont = 1
    for i=0:(1<<(2L))-1
        state_Sz = 0
        is_valid = true
        for j=0:L-1
            state_Sz += (i>>(2j))&1
            state_Sz -= (i>>(2j+1))&1
            if (i>>(2j))&1 == 1 && (i>>(2j+1))&1 == 1
                is_valid = false
            end
        end
        if state_Sz == Sz && is_valid
            states[cont] = i
            cont += 1
        end
    end
    return states
end

"""
    get_num_spin1_states(L::Int, Sz::Int)

Compute the number of states with L sites and Sz magnetization.
"""
function get_num_spin1_states(L::Int, Sz::Int)
    return length(make_spin1_LSz_basis(L, Sz))
end

"""
    make_spin1_L_basis(L::Int)

Compute states with L spins s=1.

The states are defined as in make_spin1_LSz_basis.
"""
function make_spin1_L_basis(L::Int)
    num_states = 3^L
    states = zeros(Int, num_states)
    cont = 1
    for i=0:(1<<(2L))-1
        is_valid = true
        for j=0:L-1
            if (i>>(2j))&3 == 3
                is_valid = false
            end
        end
        if is_valid
            states[cont] = i
            cont += 1
        end
    end
    return states
end


text:
module StatUtils
using SpecialFunctions
using Random
using Statistics
using Optim
import StatsBase
using Distributions
import Distributions:pdf,logpdf, mean,var, fit_mle
include("types.jl")
include("generalized_poisson.jl")
include("hist.jl")
export ZScore,Percentile, GeneralizedPoisson

"""
Compute the spearman rank correlation between r `x1` and `x2`. Also returns a two sided p-value for the signifiance of r using a permutation test
"""
function spearmanr(x1,x2,tail::Symbol=:right, RNG=MersenneTwister(rand(UInt32)))
    if !(tail in [:left, :right, :both])
        ArgumentError("tail must be one of :left, :right, :both")
    end
    cc = StatsBase.corspearman(x1,x2)
    nnl = 0
    nns = 0
    _x2 = copy(x2)
    N = 1000
    for i in 1:N
        _cc = StatsBase.corspearman(x1,shuffle!(RNG, _x2))
        if _cc > cc
            nnl +=1
        elseif _cc < cc
            nns +=1
        end
    end
    if tail == :both
        pv =  (nns + nnl)/N
    elseif tail == :left
        pv = nns/N
    else
        pv = nnl/N
    end
    cc, pv
end

"""
Use the False Discovery Rate (FDR) procedure in

Benjamini, Y., & Hochberg, Y. (1995). Controlling the false discovery rate: a practical and powerful approach to multiple testing JR Stat Soc B 57: 289–300.

to correct for the multiple comparisons represented by `pvalues` at level `0 < q < 1.0`.
"""
function fdr(pvalues::Array{Float64,1},q=0.05)
	sidx = sortperm(pvalues)
	m = length(sidx)
	vv = [1:m;]
	idx = findlast(q*vv./m .> pvalues[sidx])
	sidx[1:idx]
end

"""
Group the elements of `A` according to `grouping`, returning a dictionary of type `Dict{eltype(grouping), typeof(A)}`, where each entry (k,v) is a vector `v` with the elements of `A` corresponding to group `k`.

	function groupby{T1<:Any,T2<:Any}(A::AbstractArray{T1,1}, grouping::AbstractArray{T2,1})
"""
function groupby(A::AbstractArray{T1,1}, grouping::AbstractArray{T2,1}) where T1<:Any where T2<:Any
	groups = Dict{T2, Array{T1,1}}()
	for (a,g) in zip(A,grouping)
		if !(g in keys(groups))
			groups[g] = T1[]
		end
		push!(groups[g], a)
	end
	groups
end

function bootstrap_median(x1::AbstractArray{T,2}, n=10_000, RNG=MersenneTwister(rand(UInt32)),nx=size(x1,2)) where T<: Real
    nbins, ntrials = size(x1)
    μ = zeros(nbins)
    σ² = fill!(similar(μ), 0.0)
    nx = min(nx, ntrials)
    xs = zeros(size(x1,1), nx)
    for i in 1:n
        for k in 1:nx
            idx1 = rand(RNG, 1:ntrials)
            for j in 1:nbins
                xs[j,k] = x1[j,idx1]
            end
        end
        for j in 1:nbins
            m = median(xs[j,:])
            μ[j] += m
            σ²[j] += m*m
        end
    end
    μ ./= n
    σ² ./= n
    σ = sqrt.(σ² - μ.*μ)
    μ, σ
end

"""
Bootstrap linear regression of `y` on `x` by performing regression `n` times on random samples from `x` and `y`.
"""
function bootstrap_regression(x::AbstractVector{T},y::AbstractVector{T},n=1000;RNG=MersenneTwister(rand(UInt32))) where T <: Real
    nx = length(x)
    ym = Dict{T,T}()
    ys = Dict{T,T}()
    nn = Dict{T,Int64}()
    for i in 1:n
        _idx = rand(RNG, 1:nx, nx)
        _x = x[_idx]
        a,b = hcat(fill!(similar(_x), 1.0), _x)\(y[_idx])
        ye = a .+ b*_x
        for (_xe, _ye) in zip(_x, ye)
            ym[_xe] = get(ym, _xe, 0.0) + _ye
            ys[_xe] = get(ys, _xe, 0.0) + _ye*_ye
            nn[_xe]  = get(nn, _xe, 0) + 1
        end
    end
    xx = sort(collect(keys(ym)))
    μ = [ym[_xe]/nn[_xe] for _xe in xx]
    σ = sqrt.([ys[_xe]/nn[_xe] - (ym[_xe]/nn[_xe])^2 for _xe in xx])
    μ, σ, xx
end

"""
Regress `y` onto `x` by using the L₁ norm
"""
function robust_regression(x::AbstractVector{T}, y::AbstractVector{T},β0=rand(Float64,2)) where T <: Real
    xt = [fill!(similar(x), one(T)) x]
    func(β) = sum(abs, y .- xt*β)
    q = optimize(func, β0)
    q.minimizer
end

"""
Compute the mean of each group as indicated by `grouping`.
"""
function Statistics.mean(x::AbstractArray{T}, grouping::AbstractVector{T2}) where T <: Real where T2 <: Integer
    ng = length(grouping)
    sx = size(x)
    dim = findfirst(sx.==ng)
    spre = sx[1:dim-1]
    spost = sx[dim+1:end]
    if dim == nothing
        error("`grouping` is not compatible with any dimension of `x`")
    end
    _grouping = compress(grouping)
    ngroups = maximum(_grouping)
    outsize = (sx[1:dim-1]..., ngroups, sx[dim+1:end]...)
    μ = fill(0.0,outsize)
    counts = fill(0.0, outsize)
    for Ipost in CartesianIndices(spost)
        for i in axes(x,dim)
            for Ipre in CartesianIndices(spre)
                μ[Ipre,_grouping[i], Ipost] += x[Ipre, i, Ipost]
                counts[Ipre, _grouping[i], Ipost] += 1.0
            end
        end
    end
    μ ./= counts
    μ
end

"""
Compress the values in `x` by finding the smallest range that encompases all of `x` and repacing each value in `x` by its index into this range.
"""
function compress(x::AbstractArray{T}) where T <: Integer
    groups = unique(x)
    sort!(groups)
    y = fill!(similar(x), zero(T))
    for (i,_x) in enumerate(x)
        y[i] = findfirst(g->_x==g, groups)
    end
    y
end

end


text:
#Imports

import Pkg; Pkg.add("WAV"); Pkg.add("Plots"); Pkg.add("TickTock"); Pkg.add("MTH229"); Pkg.add("SymPy")
using Pkg

using WAV

whiteNoise = (rand(48000)*2).-1
WAV.wavwrite(whiteNoise, "whiteNoise.wav", Fs=4800) #sample freq is 4800Hz

#white noise generator script:

function createwhiten(n)
    whiten = Vector{Float64}();
    for i in 1:n*48000
        append!(whiten, (rand()*2. - 1));
    end
    return whiten
end

noise = createwhiten(1);
WAV.wavwrite(noise, "white_noise_sound2.wav", Fs=4800);
print("Number of samples: ", length(noise))

using Plots
h = Plots.histogram(noise)
Plots.display(h)
readline() #this will stop the program at this point till you press enter

using TickTock

# comparing the performance of the two approaches to generate white noise.

# Measuring excecution time of rand() generating 10 seconds of white noise:
tick()
whiteNoise = (rand(48000)*2).-1
#WAV.wavwrite(whiteNoise, "whiteNoise_test1_rand().wav", Fs=4800) #sample freq is 4800Hz
#tok()
#x = tok()
print("Time taken for rand() function to execute generating 10 seconds of white noise: ", tok(), "\n")

# Measuring excecution time of rand() generating 20 seconds of white noise:
tick()
whiteNoise = (rand(96000)*2).-1
#WAV.wavwrite(whiteNoise, "whiteNoise_test1_rand().wav", Fs=4800) #sample freq is 4800Hz
#tok()
#x = tok()
print("Time taken for rand() function to execute generating 20 seconds of white noise: ", tok(), "\n")

# Measuring excecution time of rand() generating 30 seconds of white noise:
tick()
whiteNoise = (rand(180000)*2).-1
#WAV.wavwrite(whiteNoise, "whiteNoise_test1_rand().wav", Fs=4800) #sample freq is 4800Hz
#tok()
#x = tok()
print("Time taken for rand() function to execute generating 30 seconds of white noise: ", tok(), "\n")

# Measuring the execution time of createwhiten(n) for n = 10;
tick()
noise = createwhiten(10);
#WAV.wavwrite(noise, "white_noise_sound2_test2_createwhiten.wav", Fs=4800);
#tok()
#y = tok()
print("Time taken for createwhiten(10) function to execute: ", tok(), "\n")

# Measuring the execution time of createwhiten(n) for n = 20;
tick()
noise = createwhiten(20);
#WAV.wavwrite(noise, "white_noise_sound2_test3_createwhiten.wav", Fs=4800);
print("Time taken for createwhiten(20) function to execute: ", tok(), "\n")

# Measuring the execution time of createwhiten(n) for n = 30;
tick()
noise = createwhiten(30);
#WAV.wavwrite(noise, "white_noise_sound2_test4_createwhiten.wav", Fs=4800);
print("Time taken for createwhiten(30) function to execute: ", tok(), "\n")

# Implement the Pearson’s correlation formula a new function call it corr().

function corr(x = Vector{Float64}(), y = Vector{Float}())

    #Setting up summation of xy:
    xyVal = 0
    for i in 1:length(x)
        #xyVal = 0;
        for z in 1:length(y)
            xyVal = xyVal + x[z]*y[z]
        end
        #return xyVal
    end

    #Setting up summation of x values:
    for n in 1:length(x)
        xVal = 0;
        for j in 1:length(x)
            xVal = xVal + x[j]
        end
        #return xVal
    end

    #Setting up the summation of y values:
    for c in 1:length(y)
        yVal = 0;
        for p in 1:length(y)
            yVal = yVal + y[p]
        end
        #return yVal
    end

    #Setting up summation of x^2
    x2Val = 0
    for q in 1:length(x)
        #x2Val = 0;
        for g in 1:length(x)
            x2Val = x2Val + x[g]*x[g]
        end
        #return x2Val
    end

    #Setting up summation of y^2
    y2Val = 0
    for h in 1:length(y)
        #y2Val = 0;
        for w in 1:length(y)
            y2Val = y2Val + y[w]*y[w]
        end
        #return y2Val
    end

    r = xyVal/(sqrt(x2Val)*sqrt(y2Val));
    return r

end

corr([1,2,3], [4,5,6])

# Using your output from the createWhiten(numberOfSeconds) function compare its
# correlation against itself using your correlation function and the Statistics.cor function.

using Statistics

#Correlation of createwhiten() against itself using the corr() function created.
corr(createwhiten(1), createwhiten(1))

#Correlation of creatwhiten() against itself using the Statistics.cor function
cor(createwhiten(1), createwhiten(1))

noise1 = createwhiten(1);
print("Number of samples: ", length((rand(48000)*2).-1), "\n")
print("Number of samples: ", length(noise1), "\n")
cor((rand(48000)*2).-1, createwhiten(1))

using Base
using MTH229
using Plots
using SymPy

#defining  variables used
#time
t=0:0.01:1 #100 samples
u=0:0.005:1 #200 samples used to plot and
#angular frequency
freq1=20*pi
freq2=50*2pi
freq3=2*pi
#empty array for correlation coefficients
k=zeros(0)

#Correlation coefficient vs shift, f = 20 Hz
for i in 0:0.005:1
    f(t)=sin.(freq1*(t))
    #shifted function
    g(t)=sin.(freq1*(t.-i))
    push!(k,cor(g(t),f(t)))
end

plot(u,k,
title="Corelation Coefficient vs shift, f= 20 Hz",
label=false,
lw=2,
xlabel= "Shift",
ylabel= "Correlation Coefficient")

#plot of graph and shifted graphs from above
f(t)= sin.(freq1*(t))
g(t)=sin.(freq1*(t).-pi/2)
i(t)=sin.(freq1*(t).-pi)
plot(t,f(t),
title="Graph displaying normal and shifted plots",
label=true,
lw=2,
xlabel= "t",
ylabel= "Ampltude",
)
plot!(t,g(t))
plot!(t,i(t))

#Correlation coefficient vs shift, f = 2 Hz
h=zeros(0)
for i in 0:0.005:1
    f(t)=sin.(freq3*(t))
    #shifted function
    g(t)=sin.(freq3*(t.-i))
    push!(h,cor(g(t),f(t)))
end
plot(u,h ,
title="Corelation Coefficient vs shift, f= 2 Hz",
label=false,
lw=2,
xlabel= "Shift in time",
ylabel= "Correlation Coefficient")

#Correlation coefficient vs shift, f = 50 Hz
l=zeros(0)
t=0:0.001:1 #1000 samples
u=0:0.001:1 #200 samples used to plot and
for i in 0:0.001:1
    f(t)=sin.(freq2*(t))
    #shifted function
    g(t)=sin.(freq2*(t.-i))
    push!(l,cor(g(t),f(t)))
end
plot(u,l ,
title="Corelation Coefficient vs shift, f= 50 Hz",
label=false,
lw=2,
xlabel= "Shift in time",
ylabel= "Correlation Coefficient")

max_stars_repo_path: Lab 1 EEE4120F (1).jl

max_issues_repo_path: Lab 1 EEE4120F (1).jl

max_forks_repo_path: Lab 1 EEE4120F (1).jl


text:
using Distributions, Zygote

@testset "gaussian" begin
    # @testset "Gaussian" begin
    #     rng = MersenneTwister(123456)
    #     D = 11
    #     S = 7
    #     m = randn(rng, D)
    #     A = randn(rng, D, D)
    #     Σ = Symmetric(A * A' + I)

    #     distributions_gaussian = MvNormal(m, Σ)
    #     my_gaussian = Gaussian(m, cholesky(Σ))

    #     X = rand(rng, distributions_gaussian, S)
    #     @test logpdf(distributions_gaussian, X) ≈ logpdf(my_gaussian, X)
    # end
    # @testset "normal_logpdf" begin
    #     rng = MersenneTwister(123456)
    #     P, Q = 128, 256
    #     Y = randn(rng, P, Q)
    #     M = randn(rng, P, Q)
    #     σs = exp.(randn(rng, P, Q)) .+ 1e-6
    #     @test normal_logpdf(Y, M, σs) ≈ logpdf.(Normal.(M, σs), Y)
    # end
    # @testset "StandardGaussian" begin
    #     rng = MersenneTwister(123456)
    #     D = 3

    #     X = StandardGaussian(D)
    #     @test length(X) == 3

    #     # Ensure the logpdf is the same as a regular old Gaussian.
    #     X_regular = Gaussian(zeros(D), cholesky(Matrix{Float64}(I, D, D)))

    #     x = rand(rng, X_regular)
    #     @test logpdf(X, x) ≈ logpdf(X_regular, x)

    #     xs = rand(rng, X_regular, 5)
    #     @test logpdf(X, xs) ≈ logpdf(X_regular, xs)

    #     # Ensure that sampling behaviour is also roughly the same.
    #     @test rand(MersenneTwister(123456), X) ≈ rand(MersenneTwister(123456), X_regular)
    #     @test rand(MersenneTwister(123456), X, 5) ≈ rand(MersenneTwister(123456), X_regular, 5)
    # end
    # @testset "GaussianPair" begin
    #     rng = MersenneTwister(123456)

    #     # First Gaussian.
    #     D1 = 11
    #     m1 = randn(rng, D1)
    #     A1 = UpperTriangular(randn(rng, D1, D1))
    #     Σ1 = Cholesky(to_positive_definite_softplus(A1), 'U', 0)
    #     X1 = Gaussian(m1, Σ1)

    #     # Second Gaussian.
    #     D2 = 7
    #     m2 = randn(rng, D2)
    #     A2 = UpperTriangular(randn(rng, D2, D2))
    #     Σ2 = Cholesky(to_positive_definite_softplus(A2), 'U', 0)
    #     X2 = Gaussian(m2, Σ2)

    #     # Create a Gaussian pair.
    #     X = GaussianPair(X1, X2)

    #     # Check that rand works.
    #     S = 999
    #     y = rand(rng, X)
    #     Y = rand(rng, X, S)

    #     # Check that logpdf is consistent with independent Gaussians.
    #     @test length(X) == length(X1) + length(X2)
    #     @test logpdf(X, y) ≈ logpdf(X.x1, y.x1) + logpdf(X.x2, y.x2)
    #     @test logpdf(X, Y) ≈ logpdf(X.x1, Y.x1) + logpdf(X.x2, Y.x2)
    # end
    @testset "natural parameters" begin
        rng = MersenneTwister(123456)
        D = 3
        m = randn(rng, D)
        A = randn(rng, D, D)
        S = Symmetric(A'A + I)

        θ₁, θ₂ = standard_to_natural(m, S)
        m′, S′ = natural_to_standard(θ₁, θ₂)

        @test m ≈ m′
        @test S ≈ S′

        Δθ₁ = randn(rng, D)

        A_ = randn(rng, D, D)
        Δθ₂ = Symmetric(A_ * A_' + I)
        frule_test(natural_to_standard, (θ₁, Δθ₁), (θ₂, Δθ₂); fdm=backward_fdm(5, 1))
    end
    @testset "expectation parameters" begin
        rng = MersenneTwister(123456)
        D = 3
        m = randn(rng, D)
        A = randn(rng, D, D)
        S = Symmetric(A'A + I)

        θ₁, θ₂ = standard_to_expectation(m, S)
        m′, S′ = expectation_to_standard(θ₁, θ₂)

        @test m ≈ m′
        @test S ≈ S′

        Δm = randn(rng, D)
        ΔS_root = randn(rng, D, D)
        ΔS = Symmetric(ΔS_root * ΔS_root' + I)
        frule_test(standard_to_expectation, (m, Δm), (S, ΔS); fdm=forward_fdm(5, 1))
    end
    @testset "expectation / natural parameters conversions" begin
        rng = MersenneTwister(123456)
        D = 3
        m = randn(rng, D)
        A = randn(rng, D, D)
        S = Symmetric(A'A + I)

        @testset "natural to expectation" begin
            θ₁, θ₂ = standard_to_natural(m, S)
            η₁, η₂ = natural_to_expectation(θ₁, θ₂)
            m′, S′ = expectation_to_standard(η₁, η₂)
            @test m ≈ m′
            @test S ≈ S′

            Δθ₁ = randn(rng, D)
            Δθ₂_root = randn(rng, D, D)
            Δθ₂ = Symmetric(Δθ₂_root * Δθ₂_root' + I)
            frule_test(natural_to_expectation, (θ₁, Δθ₁), (θ₂, Δθ₂); fdm=backward_fdm(5, 1))
        end

        @testset "expectation to natural" begin
            η₁, η₂ = standard_to_expectation(m, S)
            θ₁, θ₂ = expectation_to_natural(η₁, η₂)
            m′, S′ = natural_to_standard(θ₁, θ₂)
            @test m ≈ m′
            @test S ≈ S′
        end
    end
    @testset "unconstrained" begin
        rng = MersenneTwister(123456)
        D = 3
        m = randn(rng, D)
        U = UpperTriangular(randn(rng, D, D))

        @testset "standard is psd" begin
            _, S = unconstrained_to_standard(m, U)
            @test all(eigvals(S) .> 0)
        end
        @testset "standard_to_unconstrained" begin

            # Ensure inverse is correct.
            ms, Ss = unconstrained_to_standard(m, U)
            m′, U′ = standard_to_unconstrained(ms, Ss)
            @test m ≈ m′
            @test U ≈ U′

            # Use positive definite forwards sensitivity to test Ss.
            Δms = randn(rng, length(ms))
            _ΔSs = randn(rng, size(Ss))
            ΔSs = Symmetric(_ΔSs * _ΔSs' + I)

            fdm = forward_fdm(5, 1)
            frule_test(standard_to_unconstrained, (ms, Δms), (Ss, ΔSs); fdm=fdm)
        end
        @testset "natural is nsd" begin
            θ₁, θ₂ = unconstrained_to_natural(m, U)
            @test all(eigvals(θ₂) .< 0)
        end
        @testset "natural_to_unconstrained" begin
            m′, U′ = natural_to_unconstrained(unconstrained_to_natural(m, U)...)
            @test m ≈ m′
            @test U ≈ U′

            θ₁, θ₂ = unconstrained_to_natural(m, U)
            θ₁_comp, θ₂_comp = standard_to_natural(unconstrained_to_standard(m, U)...)
            @test θ₁_comp ≈ θ₁
            @test θ₂_comp ≈ θ₂

            Δθ₁ = randn(rng, length(θ₁))
            Δθ₂_ = randn(rng, size(θ₂))
            Δθ₂ = Symmetric(Δθ₂_ * Δθ₂_')
            fdm = backward_fdm(5, 1)
            frule_test(natural_to_unconstrained, (θ₁, Δθ₁), (θ₂, Δθ₂); fdm=fdm)
        end
        @testset "expectation is psd" begin
            _, S = unconstrained_to_expectation(m, U)
            @test all(eigvals(S) .> 0)
        end
        @testset "expectation_to_unconstrained" begin
            m′, U′ = expectation_to_unconstrained(unconstrained_to_expectation(m, U)...)
            @test m ≈ m′
            @test U ≈ U′

            η₁, η₂ = unconstrained_to_expectation(m, U)
            η₁_comp, η₂_comp = standard_to_expectation(unconstrained_to_standard(m, U)...)
            @test η₁ ≈ η₁_comp
            @test η₂ ≈ η₂_comp
        end
        # @testset "natural gradient tests" begin
        #     rng = MersenneTwister(123456)
        #     D = 3
        #     η̄₁, η̄₂ = randn(rng, D), randn(rng, D, D)
        #     η̄ = hcat(η̄₁, η̄₂)

        #     m, U = randn(rng, D), UpperTriangular(randn(rng, D, D))
        #     _, S = unconstrained_to_standard(m, U)
        #     θ₁, θ₂ = unconstrained_to_natural(m, U)

        #     adjoint_test(
        #         (θ₁, θ₂) -> hcat(natural_to_expectation(θ₁, θ₂)...),
        #         η̄, θ₁, θ₂,
        #     )

        #     adjoint_test(
        #         (m, S) -> hcat(standard_to_expectation(m, S)...),
        #         η̄, m, S,
        #     )

        #     adjoint_test(
        #         (m, U) -> hcat(unconstrained_to_standard(m, UpperTriangular(U))...),
        #         η̄, m, U,
        #     )

        #     adjoint_test(
        #         U -> Matrix(UpperTriangular(U)'UpperTriangular(U)),
        #         randn(rng, D, D), randn(rng, D, D),
        #     )

        #     adjoint_test(
        #         (m, U) -> hcat(unconstrained_to_expectation(m, UpperTriangular(U))...),
        #         η̄, m, U,
        #     )

        #     # Compute sensible norm in the first way.
        #     _, back_nat_to_exp = Zygote.pullback(natural_to_expectation, θ₁, θ₂)
        #     θ̄₁, θ̄₂ = back_nat_to_exp((η̄₁, η̄₂))
        #     display(hcat(θ̄₁, θ̄₂))
        #     println()

        #     norm_1 = dot(θ̄₁, η̄₁) + dot(θ̄₂, η̄₂)
        #     @show dot(θ̄₁, η̄₁), dot(θ̄₂, η̄₂)
        #     @show norm_1

        #     # Compute sensible norm in the second way.
        #     _, back_unconstrained_to_expectation = Zygote.pullback(
        #         unconstrained_to_expectation, m, U,
        #     )
        #     (m̄, Ū) = back_unconstrained_to_expectation((η̄₁, η̄₂))

        #     _, pushforward_nat_to_uncon = frule(natural_to_unconstrained, θ₁, θ₂)
        #     Δm_nat, ΔU_nat = pushforward_nat_to_uncon(DoesNotExist(), η̄₁, η̄₂)

        #     @show dot(m̄, Δm_nat), dot(Ū, ΔU_nat)
        #     norm_2 = dot(m̄, Δm_nat) + dot(Ū, ΔU_nat)
        #     @show norm_2
        # end

        @testset "positive definite tests" begin
            rng = MersenneTwister(123456)
            D = 2

            function expectation_to_natural_vec(η)
                η₁ = η[1:D]
                η₂ = Symmetric(reshape(η[D+1:end], D, D))
                θ₁, θ₂ = expectation_to_natural(η₁, η₂)
                return vec(hcat(θ₁, θ₂))
            end

            function natural_to_expectation_vec(θ)
                θ₁ = θ[1:D]
                θ₂ = Symmetric(reshape(θ[D+1:end], D, D))
                η₁, η₂ = natural_to_expectation(θ₁, θ₂)
                return vec(hcat(η₁, η₂))
            end

            # Construct natural parameters and convert to vector.
            m = randn(rng, D)
            U = UpperTriangular(randn(rng, D, D))
            θ₁, θ₂ = unconstrained_to_natural(m, U)
            θ = vec(hcat(θ₁, θ₂))
            η = natural_to_expectation_vec(θ)

            m, S = unconstrained_to_standard(m, U)
            θ_std = vec(hcat(m, S))

            println("vals")
            display(θ₁)
            println()
            display(θ₂)
            println()

            function expectation_mat_vec(θ₂)
                _, η₂ = natural_to_expectation(zeros(D), reshape(θ₂, D, D))
                return vec(η₂)
            end

            function natural_to_standard_vec(θ)
                m = θ[1:D]
                S = reshape(θ[D+1:end], D, D)
                η₁, η₂ = Zygote.@showgrad(natural_to_standard(m, S))
                return vec(hcat(η₁, η₂))
            end

            function standard_to_expectation_vec(θ)
                θ₁ = θ[1:D]
                θ₂ = reshape(θ[D+1:end], D, D)
                η₁, η₂ = Zygote.@showgrad(standard_to_expectation(θ₁, θ₂))
                return vec(hcat(η₁, η₂))
            end

            J = Matrix{Float64}(undef, length(θ), length(θ))
            Jinv = Matrix{Float64}(undef, length(θ), length(θ))
            Jforward = Matrix{Float64}(undef, length(θ), length(θ))
            for d in eachindex(θ)

                println("d = $d")

                # Define dth seed for reverse-pass.
                Δη = zeros(length(θ))
                Δη[d] = 1.0

                # Compute dth column of J
                _, back = Zygote.pullback(natural_to_expectation_vec, θ)
                Δθ = first(back(Δη))
                J[:, d] .= Δθ

                # Compute dth column of Jinv
                _, back = Zygote.pullback(expectation_to_natural_vec, η)
                Jinv[:, d] = first(back(Δη))

                # # Compute dth column of Jforward
                # Δθ = Δη
                # Δθ₁ = Δθ[1:D]
                # Δθ₂ = reshape(Δθ[D+1:end], D, D)
                # @show Δθ₁
                # @show Δθ₂
                # _, pf = frule(natural_to_expectation, θ₁, θ₂)
                # Δη₁, Δη₂ = pf(DoesNotExist(), Δθ₁, Δθ₂)
                # Jforward[1:D, d] = Δη₁
                # Jforward[D+1:end, d] = vec(Δη₂)

                println()
                println()
                println()
            end

            println("J'")
            display(J)
            println()
            inds = vcat(1:3, 5:6)
            display(J[inds, inds])
            println()

            display(J - J')
            println()

            println("Jinv")
            display(Jinv)
            println()

            display(J[inds, inds] - inv(Jinv[inds, inds]))
            println()
            @show eigvals(J[inds, inds])

            # println("Jforward")
            # display(Jforward)
            # println()

            # display(Jforward - Jforward')
            # println()

            # display(eigvals(Jforward))
            # println()

            # display(Jforward - J')
            # println()

            v = randn(length(θ))
            _, back = Zygote.pullback(natural_to_expectation_vec, θ)
            @test first(back(v))'v ≈ v' * (J * v)

            @test J ≈ J'
            @test all(eigvals(Symmetric(J)) .> 0)

            # J_mat = Matrix{Float64}(undef, length(θ₂), length(θ₂))
            # for d in eachindex(θ₂)

            #     # Define dth seed for reverse-pass.
            #     Δη = zeros(length(θ₂))
            #     Δη[d] = 1.0

            #     # Compute dth column of J
            #     _, back = Zygote.pullback(expectation_mat_vec, vec(θ₂))
            #     Δθ₂ = first(back(Δη))
            #     J_mat[:, d] .= Δθ₂
            # end

            # @show "J_mat stuff"

            # display(J_mat)
            # println()

            # display(J_mat - J_mat')
            # println()

            # Check that stuff composes.

            # J_stn = Matrix{Float64}(undef, length(θ), length(θ))
            # for d in eachindex(θ)

            #     println("d = $d")

            #     # Define dth seed for reverse-pass.
            #     Δη = zeros(length(θ))
            #     Δη[d] = 1.0

            #     # Compute dth column of J
            #     _, back = Zygote.pullback(natural_to_standard_vec, θ)
            #     Δθ = first(back(Δη))
            #     J_stn[:, d] .= Δθ

            #     println()
            #     println()
            #     println()
            # end

            # J_exp = Matrix{Float64}(undef, length(θ), length(θ))
            # for d in eachindex(θ)

            #     println("d = $d")

            #     # Define dth seed for reverse-pass.
            #     Δη = zeros(length(θ))
            #     Δη[d] = 1.0

            #     # Compute dth column of J
            #     _, back = Zygote.pullback(standard_to_expectation_vec, θ_std)
            #     Δθ = first(back(Δη))
            #     J_exp[:, d] .= Δθ

            #     println()
            #     println()
            #     println()
            # end

            # display(J_exp')
            # println()

            # display(θ₁)
            # println()
            # display(-inv(θ₂) ./ 2)
            # println()

            # display(J_stn')
            # println()

            # for d in eachindex(θ₂)

            #     println("d = $d")
            #     Δη = zeros(length(θ₂))
            #     Δη[d] = 1.0
            #     Δη = reshape(Δη, D, D)
            #     display(- inv(θ₂) * Δη * inv(θ₂) ./ 4)
            #     println()
            # end

            # display(J_exp' * J_stn')
            # println()

            # display(J_exp' * J_stn' - J')
            # println()
        end
    end
end


text:
# TODO criar uma função annual_costs mais geral, por exemplo para os NonDispatchables, DieselGenerator e Battery
function annual_costs(nd::NonDispatchables, mg::Microgrid)

    # discount factor for each year of the project
    discount_factors = [ 1/((1 + mg.project.discount_rate)^i) for i=1:mg.project.lifetime ]

    # number of replacements
    replacements_number = ceil(Integer, mg.project.lifetime/nd.lifetime) - 1
    # years that the replacements happen
    replacement_years = [i*nd.lifetime for i=1:replacements_number]
    # discount factors for the replacements years
    replacement_factors = [1/(1 + mg.project.discount_rate)^i for i in replacement_years]

    # component remaining life at the project end
    remaining_life = nd.lifetime - (mg.project.lifetime - nd.lifetime * replacements_number)
    # proportional unitary salvage cost
    proportional_salvage_cost = nd.salvage_cost * remaining_life / nd.lifetime

    # present investment cost
    investment_cost = nd.investment_cost * nd.power_rated
    # present operation and maintenance cost
    om_cost = sum(nd.om_cost * nd.power_rated * discount_factors)
    # present replacement cost
    if replacements_number == 0
        replacement_cost = 0
    else
        replacement_cost = sum(nd.replacement_cost * nd.power_rated * replacement_factors)
    end
    # present salvage cost
    if remaining_life == 0
        salvage_cost = 0
    else
        salvage_cost = proportional_salvage_cost * nd.power_rated * discount_factors[mg.project.lifetime]
    end

    total_cost = investment_cost + replacement_cost + om_cost - salvage_cost

    return [total_cost, investment_cost, om_cost, replacement_cost, salvage_cost]
end

function annual_costs(dg::DieselGenerator, mg::Microgrid, opervarsaggr::OperVarsAggr)

    # discount factor for each year of the project
    discount_factors = [ 1/((1 + mg.project.discount_rate)^i) for i=1:mg.project.lifetime ]

    # total diesel generator operation hours over the project lifetime
    total_DG_operation_hours = mg.project.lifetime * opervarsaggr.DG_operation_hours

    # number of replacements
    replacements_number = ceil(Integer, total_DG_operation_hours/dg.lifetime) - 1
    # years that the replacements happen
    replacement_years = [i*(dg.lifetime/opervarsaggr.DG_operation_hours) for i=1:replacements_number]     # TODO verify
    # discount factors for the replacements years
    replacement_factors = [1/(1 + mg.project.discount_rate)^i for i in replacement_years]

    # component remaining life at the project end
    remaining_life = dg.lifetime - (total_DG_operation_hours - dg.lifetime * replacements_number)
    # proportional unitary salvage cost
    proportional_salvage_cost = dg.salvage_cost * remaining_life / dg.lifetime

    # present investment cost
    investment_cost = dg.investment_cost * dg.power_rated
    # present operation and maintenance cost
    om_cost = sum(dg.om_cost * dg.power_rated * opervarsaggr.DG_operation_hours * discount_factors) # depends on the nb of the DG working Hours
    # present replacement cost
    if replacements_number == 0
        replacement_cost = 0
    else
        replacement_cost = sum(dg.replacement_cost * dg.power_rated * replacement_factors)
    end
    # present salvage cost
    if remaining_life == 0
        salvage_cost = 0
    else
        salvage_cost = proportional_salvage_cost * dg.power_rated * discount_factors[mg.project.lifetime]
    end

    fuel_cost = sum(dg.fuel_cost * opervarsaggr.fuel_consumption * discount_factors)

    total_cost = investment_cost + replacement_cost + om_cost - salvage_cost + fuel_cost

    return [total_cost, investment_cost, om_cost, replacement_cost, salvage_cost, fuel_cost]
end

function annual_costs(bt::Battery, mg::Microgrid, opervarsaggr::OperVarsAggr)

    # discount factor for each year of the project
    discount_factors = [ 1/((1 + mg.project.discount_rate)^i) for i=1:mg.project.lifetime ]

    # minimum battery lifetime between years lifetime and number of cycles lifetime
    if opervarsaggr.annual_throughput == 0
        bt_lifetime_min = bt.lifetime
    else
        bt_lifetime_min = minimum( [(bt.energy_max * bt.lifetime_throughput/opervarsaggr.annual_throughput), bt.lifetime] )
    end

    # number of replacements
    replacements_number = ceil(Integer, mg.project.lifetime/bt_lifetime_min) - 1
    # years that the replacements happen
    replacement_years = [i*bt_lifetime_min for i=1:replacements_number]
    # discount factors for the replacements years
    replacement_factors = [1/(1 + mg.project.discount_rate)^i for i in replacement_years]

    # component remaining life at the project end
    remaining_life = bt_lifetime_min - (mg.project.lifetime - bt_lifetime_min* replacements_number)
    # proportional unitary salvage cost
    proportional_salvage_cost = bt.salvage_cost * remaining_life / bt_lifetime_min

    # present investment cost
    investment_cost = bt.investment_cost * bt.energy_max
    # present operation and maintenance cost
    om_cost = sum(bt.om_cost * bt.energy_max * discount_factors)
    # present replacement cost
    if replacements_number == 0
        replacement_cost = 0
    else
        replacement_cost = sum(bt.replacement_cost * bt.energy_max * replacement_factors)
    end
    # present salvage cost
    if remaining_life == 0
        salvage_cost = 0
    else
        salvage_cost = proportional_salvage_cost * bt.energy_max * discount_factors[mg.project.lifetime]
    end

    total_cost = investment_cost + replacement_cost + om_cost - salvage_cost

    return [total_cost, investment_cost, om_cost, replacement_cost, salvage_cost]
end

"""
    economics(mg::Microgrid, opervarsaggr::OperVarsAggr)

Return the economics results for the microgrid `mg` and
the aggregated operation variables `OperVarsAggr`.

See also: [`aggregation`](@ref)
"""
function economics(mg::Microgrid, opervarsaggr::OperVarsAggr)

    # discount factor for each year of the project
    discount_factors = [ 1/((1 + mg.project.discount_rate)^i) for i=1:mg.project.lifetime ]

    # Photovoltaic costs initialization
    PV_total_cost = 0.
    PV_investment_cost = 0.
    PV_om_cost = 0.
    PV_replacement_cost= 0.
    PV_salvage_cost = 0.
    # Wind power costs initialization
    WT_total_cost = 0.
    WT_investment_cost = 0.
    WT_om_cost = 0.
    WT_replacement_cost= 0.
    WT_salvage_cost = 0.
    # Diesel generator costs initialization
    #= DG_total_cost = 0.
    DG_investment_cost = 0.
    DG_om_cost = 0.
    DG_replacement_cost= 0.
    DG_salvage_cost = 0.
    DG_fuel_cost = 0. =#

    # NonDispatchables costs
    for i=1:length(mg.nondispatchables)
        if typeof(mg.nondispatchables[i]) == Photovoltaic
            PV_total_cost, PV_investment_cost, PV_om_cost, PV_replacement_cost, PV_salvage_cost = annual_costs(mg.nondispatchables[i], mg)
        elseif typeof(mg.nondispatchables[i]) == WindPower
            WT_total_cost, WT_investment_cost, WT_om_cost, WT_replacement_cost, WT_salvage_cost = annual_costs(mg.nondispatchables[i], mg)
        end
    end

    # DieselGenerator costs
    DG_total_cost, DG_investment_cost, DG_om_cost, DG_replacement_cost, DG_salvage_cost, DG_fuel_cost = annual_costs(mg.dieselgenerator, mg, opervarsaggr)

    # Battery costs
    BT_total_cost, BT_investment_cost, BT_om_cost, BT_replacement_cost, BT_salvage_cost = annual_costs(mg.battery, mg, opervarsaggr)

    # SUMMARY
    # total present investment cost
    total_investment_cost = DG_investment_cost + BT_investment_cost + PV_investment_cost + WT_investment_cost
    # total present replacement cost
    total_replacement_cost = DG_replacement_cost + BT_replacement_cost + PV_replacement_cost + WT_replacement_cost
    # total present operation and maintenance cost
    total_om_cost = DG_om_cost + BT_om_cost + PV_om_cost + WT_om_cost
    # total present salvage cost
    total_salvage_cost = DG_salvage_cost + BT_salvage_cost + PV_salvage_cost + WT_salvage_cost
    # net present cost
    npc = DG_total_cost + BT_total_cost + PV_total_cost + WT_total_cost

    # recovery factor
    recovery_factor = (mg.project.discount_rate * (1 + mg.project.discount_rate)^mg.project.lifetime)/((1 + mg.project.discount_rate)^mg.project.lifetime - 1)
    # total annualized cost
    annualized_cost = npc * recovery_factor
    # cost of energy
    coe = annualized_cost / opervarsaggr.energy_served

    # energy served over the project lifetime
    energy_served_lifetime = opervarsaggr.energy_served * sum([1.0; discount_factors[1:length(discount_factors)-1]])
    # levelized cost of energy
    lcoe = npc / energy_served_lifetime

    costs = TotalCosts(lcoe, coe, npc,
            total_investment_cost, total_replacement_cost, total_om_cost, total_salvage_cost,
            DG_total_cost, DG_investment_cost, DG_replacement_cost, DG_om_cost, DG_salvage_cost, DG_fuel_cost,
            BT_total_cost, BT_investment_cost, BT_replacement_cost, BT_om_cost, BT_salvage_cost,
            PV_total_cost, PV_investment_cost, PV_replacement_cost, PV_om_cost, PV_salvage_cost,
            WT_total_cost, WT_investment_cost, WT_replacement_cost, WT_om_cost, WT_salvage_cost)

    return costs
end


text:
using Sundials
using Test

# Some tests use @cfunction, but at the moment that isn't supported on non-Intel platforms
const SUPPORT_CFUNCTION = Sys.ARCH ∉ (:aarch64, :ppc64le, :powerpc64le) || startswith(lowercase(String(Sys.ARCH)), "arm")

@testset "CVODE" begin
    @testset "Roberts CVODE Simplified" begin
        include("cvode_Roberts_simplified.jl")
    end
    SUPPORT_CFUNCTION && @testset "Roberts CVODE Direct" begin
        include("cvode_Roberts_dns.jl")
    end
    #@testset "CVODES Direct" begin include("cvodes_dns.jl") end
end

@testset "IDA" begin
    @testset "Roberts IDA Simplified" begin
        include("ida_Roberts_simplified.jl")
    end
    SUPPORT_CFUNCTION && @testset "Roberts IDA Direct" begin
        include("ida_Roberts_dns.jl")
    end
    @testset "Heat IDA Direct" begin
        include("ida_Heat2D.jl")
    end
    # Commented out because still uses the syntax from Grid which is a deprecated package
    #@testset "Cable IDA Direct" begin include("ida_Cable.jl") end
end

SUPPORT_CFUNCTION && @testset "ARK" begin
    @testset "Roberts ARKStep Direct" begin
        include("arkstep_Roberts_dns.jl")
    end
    @testset "NonLinear ERKStep Direct" begin
        include("erkstep_nonlin.jl")
    end
    #@testset "MRI two way couple" begin include("mri_twowaycouple.jl") end
end

@testset "Kinsol" begin
    @testset "Kinsol Simplified" begin
        include("kinsol_mkin_simplified.jl")
    end
    SUPPORT_CFUNCTION && @testset "Kinsol MKin" begin
        include("kinsol_mkinTest.jl")
    end
    @testset "Kinsol Banded" begin
        include("kinsol_banded.jl")
    end
end
@testset "Handle Tests" begin
    include("handle_tests.jl")
end

@testset "Common Interface" begin
    @testset "CVODE" begin
        include("common_interface/cvode.jl")
    end
    @testset "ARKODE" begin
        include("common_interface/arkode.jl")
    end
    @testset "IDA" begin
        include("common_interface/ida.jl")
    end
    @testset "Jacobians" begin
        include("common_interface/jacobians.jl")
    end
    @testset "Callbacks" begin
        include("common_interface/callbacks.jl")
    end
    @testset "Iterator" begin
        include("common_interface/iterators.jl")
    end
    @testset "Errors" begin
        include("common_interface/errors.jl")
    end
    @testset "Mass Matrix" begin
        include("common_interface/mass_matrix.jl")
    end
end