text:
import Images

const writedir = joinpath(tempdir(), "Images")

if !isdir(writedir)
    mkdir(writedir)
end

a = rand(2,2)
fn = joinpath(writedir, "2by2.png")
Images.imwrite(a, fn)
b = Images.imread(fn)
@assert convert(Array, b) == scale(Images.scaleinfo(a), a)
aa = int8(127*a)
Images.imwrite(aa, fn)
b = Images.imread(fn)
@assert convert(Array{Int8}, b) == aa
aa = uint8(255*a)
Images.imwrite(aa, fn)
b = Images.imread(fn)
@assert convert(Array, b) == aa
aa = uint16(65535*a)
Images.imwrite(aa, fn)
b = Images.imread(fn)
@assert convert(Array, b) == aa


text:
# Julia 0.6.2
import Base.*
using LinearAlgebra

× = *
ᵀ = transpose
(x)ᵀ = transpose(x)
⋅ = dot
∑(x) = reduce(+, x)
√(x) = sqrt(x)

# ascii and unicode euclidean distance implementation
euc_dist₁(x, y) = sqrt(sum((x-y).^2))
euc_dist₂(x, y) = √(∑((x-y)⋅(x-y)))

a, b = rand(5), rand(5)
println(euc_dist₁(a, b) - euc_dist₂(a, b))

# ascii and unicode cosine similarity implementation
⁻¹ = inv
(x)⁻¹ = 1/x
ϕ = norm
cos_dist₁(x, y) = dot(x, y) / (norm(x) * norm(y))
cos_dist₂(x, y) = (x ⋅ y) × ((ϕ(x) × ϕ(y))⁻¹)

println(cos_dist₁(a, b) - cos_dist₂(a, b))

ext: jl

max_stars_repo_name: cdluminate/MyNotes

max_issues_repo_name: cdluminate/MyNotes

max_forks_repo_name: cdluminate/MyNotes


text:
using PkgSRA, Test, Random
Random.seed!(13)

# include("../utils/sindy_turing_utils.jl")
# include("../utils/sparse_regression_functions.jl")
# include("../utils/sindy_statistics_utils.jl")

# Generate test data: Lorenz
include("../examples/example_lorenz.jl")
dat = Array(solve_lorenz_system())
numerical_grad = numerical_derivative(dat, ts)

# Test new interface
alg = slstQuantile(4, 0.2) # Will get extra terms
sindy_library = Dict("cross_terms"=>2,"constant"=>nothing);
test_model = sindy(dat, numerical_grad, ts,
                        library=sindy_library,
                        optimizer=alg)

# Built correct terms: names
named_terms = build_term_names(test_model)
@testset "Term Metadata" begin
    @test named_terms ==
        ["x", "y", "z", "", "xx", "xy", "xz", "yy", "yz", "zz"]
end

# Built correct terms: values
library_functions = convert_string2function(test_model.library)
terms = calc_augmented_data(dat, library_functions)

@testset "Data Augmentation" begin
    @test size(terms,1) == length(named_terms)
    @test terms[1:3,:] == dat
    @test terms[4,:] == ones(size(dat,2))
    @test terms[5:10,:] == calc_cross_terms(dat, 2)
end;

# Nonzero term and type tests
t1 = get_nonzero_terms(test_model)
t2 = get_nonzero_terms(core_dyn_true)
# typeof(test_model) <: sindycModel
# @testset "Type Tests" begin
@test typeof(test_model) <: sindyModel
@test issubset(t2, t1)
# end

###
### Cross validation functions
###

# First, retraining functions
test_model2 = sindy_retrain(test_model, dat, numerical_grad)
t1 = get_nonzero_terms(test_model)
t2 = get_nonzero_terms(test_model2)

@testset "Retrained model" begin
    @test typeof(test_model2) <: sindyModel
    @test t1==t2
    @test calc_coefficient_error(test_model, test_model2) ≈ 0.0
end

# Also test cross validation function
null_model = sindy(dat, numerical_grad, ts,
                        library=sindy_library,
                        optimizer=slstHard(100)) # i.e. not sparse
err_test = sindy_cross_validate(test_model, dat, numerical_grad)
err_null = sindy_cross_validate(null_model, dat, numerical_grad)

@testset "Cross Validation" begin
    @test err_test < err_null
end

## One more cross validation test
include("../examples/example_lotkaVolterra.jl")
dat = Array(solve_lv_system())
numerical_grad = numerical_derivative(dat, ts)

# First, a decent model on one variable
sindy_library = Dict("cross_terms"=>2,"constant"=>nothing);
lib = convert_string2function(sindy_library)
#      x   y   c   xx xy yy
A1 = [[0.7 0   0   0 -1.3 0 ];
      [0  0   0   0   0   1e-6 ]]
n = size(A, 1)
m1 = sindyModel(ts, A1, lib ,["x", "y"])

# Second, a decent model on both variables
#      x   y   c   xx xy yy
A2 = [[0.7 0   0   0 -1.3 1e-6 ];
      [0  -0.7 0  1e-6 0.7  0]]
m2 = sindyModel(ts, A2, lib ,["x", "y"])

# Third, another decent model on both variables
#      x   y   c   xx xy yy
A3 = [[0.7 0   0   0 -1.3 0 ];
      [1e-6 -0.9 0 0 0.8  0]]
m3 = sindyModel(ts, A3, lib ,["x", "y"])

## Calculate errors
# Coefficients
cerr1 = calc_coefficient_error(m1, core_dyn_true)
cerr2 = calc_coefficient_error(m2, core_dyn_true)
# Residual
ind = 1:500
rss0 = rss_sindy_derivs(core_dyn_true, dat[:,ind], numerical_grad[:,ind])
rss1 = rss_sindy_derivs(m1, dat[:,ind], numerical_grad[:,ind])
rss2 = rss_sindy_derivs(m2, dat[:,ind], numerical_grad[:,ind])
rss3 = rss_sindy_derivs(m3, dat[:,ind], numerical_grad[:,ind])

#
@testset "Error calculations" begin
    @test cerr1 > cerr2
    @test rss1 > rss0
    @test rss2 > rss0
    @test rss1 > rss2
end

using Plots
pyplot()
u0 = dat[:,1]
plot_sindy_model(m1, u0, which_dim=2)
    plot!(m1.ts, dat[2,:])
    title!("Model 1")
plot_sindy_model(m2, u0, which_dim=2)
    plot!(m2.ts, dat[2,:])
    title!("Model 2")

# SINDY tests: up to cubed terms

# Generate test data: Van der Pol
include("../examples/example_vanDerPol.jl")
dat = Array(solve_vdp_system())
numerical_grad = numerical_derivative(dat, ts)

sindy_library = Dict("cross_terms"=>[2,3],
                    "constant"=>nothing);
test_model = sindy(dat, numerical_grad, ts,
                        library=sindy_library,
                        optimizer=alg,
                        var_names=["x", "y"])

named_terms = build_term_names(test_model)
@testset "Cubic term names" begin
    @test named_terms ==
        ["x", "y", "", "xx", "xy", "yy", "xxx", "xxy", "xyy", "yyy"]
end


text:
using GitLinks
using Test

@testset "GitLinks.jl" begin
    include("git_utils_tests.jl")
    include("server_client_tests.jl")
end


text:
using CollegeStratBase, CollegeStratPlots
using Test

csp = CollegeStratPlots;

function fig_test_setup(figName)
    fPath = joinpath(CollegeStratBase.test_dir(), figName);
    notesPath = csp.fig_notes_path(fPath);
    isfile(fPath)  &&  rm(fPath);
    isfile(notesPath)  &&  rm(notesPath);

    plot_defaults();
    return fPath, notesPath
end

function fig_save_test()
    @testset "Save figure" begin
        fPath, notesPath = fig_test_setup("fig_save_test.pdf");
        p = csp.test_bar_graph();

        fNotes = ["This is a test figure.", "With figure notes"];
        figsave(p, fPath; figNotes = fNotes);
        @test isfile(fPath)
        @test isfile(notesPath)
    end
end

function grouped_bar_test()
    @testset "Grouped bar graph" begin
        fPath, notesPath = fig_test_setup("grouped_bar_test.pdf");
        nr = 5; nc = 6;
        dataM = (1:nr) .+ (1 : nc)';
        p = grouped_bar_graph(string.(1:nr), dataM);
        figsave(p, fPath; figNotes = ["Notes"]);
        @test isfile(fPath)
        @test isfile(notesPath)
    end
end

function histogram_test()
    @testset "Histogram" begin
        fPath, notesPath = fig_test_setup("histogram_test.pdf");
        p = histogram_plot(LinRange(-3.0, 4.0, 1000); xlabel = "x")
        figsave(p, fPath);
        @test isfile(fPath);
    end
end

function subplot_test()
    @testset "Subplots" begin
        nPlots = 3;
        pV = Vector{Any}(undef, nPlots);
        for ip = 1 : nPlots
            pV[ip] = csp.test_bar_graph();
        end
        p = sub_plots(pV; link = :all);
        fPath, notesPath = fig_test_setup("subplots_test.pdf");
        figsave(p, fPath);
    end
end

function line_plot_test()
    @testset "Line plot" begin
        fPath, notesPath = fig_test_setup("line_plot_test.pdf");
        nx = 7;
        xV = LinRange(-2.0, 1.5, nx);
        yM = LinRange(0.3, 0.5, nx) .+ LinRange(1.0, 0.5, 4)';
        p = line_plot(xV, yM; xlabel = "x");
        figsave(p, fPath);
        @test isfile(fPath);

        fPath2, _ = fig_test_setup("add_line_test.pdf");
        yV = xV .+ 0.5;
        p = line_plot(xV, yV; ylabel = "y");
        add_line!(p, xV, yV .+ 1.0);
        figsave(p, fPath2);
        @test isfile(fPath2);
    end
end

@testset "CollegeStratPlots.jl" begin
    fig_save_test();
    grouped_bar_test();
    line_plot_test();
    subplot_test();
    histogram_test();
end

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


text:
module JuliRay

using LinearAlgebra
using Printf
using Colors

include("_BasicFunctions.jl")

## Translation to POV-Ray code

"""
2.3 -> "2.3"
"""
function povray_script(x::Real)
    if isfinite(x)
        return repr(x)
        # return @sprintf "%1.24f" x
    else
        error("numerical value must be finite")
    end
end

"""
[2.3, -5.2] -> "<2.3,-5.2>"
"""
function povray_script(x::Array{T,1}) where T <: Real
    return "<"*join(povray_script.(x), ", ")*">"
end

"""
[[2.3, -5.2], [-1.1, 8.2]] -> "<2.3,-5.2>, <-1.1,8.2>"
"""
function povray_script(x::Array{Array{T,1},1}) where T <: Real
    return join(povray_script.(x), ", ")
end

function povray_script(x::Array{T,1}) where T <: Color
    return join((s->"texture{pigment{"*s*"}}").(povray_script.(x)), ", ")
end

function povray_script(x::Array{T,1}, t::Array{FT,1}) where T <: Color
    return join((s->"texture{pigment{"*s*"}}").([povray_script(x[i],t[i]) for i in 1:length(x)]), ", ")
end

# Types
abstract type JR end
abstract type Object <: JR end

include("_PrimitiveObject.jl")
include("_CSG.jl")
include("_Transformation.jl")
include("_Pigment.jl")
include("_CompoundObject.jl")
include("_ParametricObject.jl")
include("_Camera.jl")

BASE_DIR = "."

# Rendering
export render
function render(object::Object; name="new", index::Int=0, camera::Camera=LngLatCamera(), lights::Array{T,1}=Light[]) where T <: Light
    width=camera.width
    height=camera.height

    if 0 < index < 1000000
        Index = "_"*(@sprintf "%06d" index)
    elseif index == 0
        Index = ""
    else
        error("index must be non-negative and less than 1000000")
    end
    if endswith(name,".pov")
        Name = name[1:end-4]*Index
    else
        Name = name*Index
    end

    str = "#version 3.7;\nglobal_settings{assumed_gamma 1.0}\n"
    str = str*povray_script(camera)*"\n"
    str = str*(join(povray_script.(lights)))
    str = str*"background{rgb<1,1,1>}"*"\n"
    str = str*povray_script(object)*"\n"
    path_pov = BASE_DIR * "/" * Name * ".pov"
    io = open(path_pov,"w")
    write(io,str)
    close(io)

    str = "Width=$width\nHeight=$height\nAntialias=On\n"
    path_ini = BASE_DIR * "/" * "povray.ini"
    io = open(path_ini,"w")
    write(io,str)
    close(io)

    run(`povray $path_pov`)
end

end # module

max_stars_repo_name: hyrodium/JuliRay

max_issues_repo_name: hyrodium/JuliRay

max_forks_repo_name: hyrodium/JuliRay


text:
module CollabUtils2ChapterPanelTrigger

using ..Ahorn, Maple

@mapdef Trigger "CollabUtils2/ChapterPanelTrigger" ChapterPanelTrigger(x::Integer, y::Integer, width::Integer=Maple.defaultTriggerWidth, height::Integer=Maple.defaultTriggerHeight,
    map::String="Celeste/1-ForsakenCity", returnToLobbyMode::String="SetReturnToHere", allowSaving::Bool=true)

const placements = Ahorn.PlacementDict(
    "Chapter Panel Trigger (Collab Utils 2 / READ DOCS)" => Ahorn.EntityPlacement(
        ChapterPanelTrigger,
        "rectangle",
    ),
)

function Ahorn.editingOptions(trigger::ChapterPanelTrigger)
    return Dict{String, Any}(
        "returnToLobbyMode" => String["SetReturnToHere", "RemoveReturn", "DoNotChangeReturn"]
    )
end

function Ahorn.nodeLimits(trigger::ChapterPanelTrigger)
    return 0, 1
end

end


text:
using FiniteVolume
using Test

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


text:
include("header.jl")

@testset "cat" begin
    o = (:delta=>0.01,:rtol=>0.001,:atol=>0.001)
    cat1(x...)=cat(x...; dims=Val(1))
    cat2(x...)=cat(x...; dims=Val(2))
    cat31(x...)=cat(x...; dims=Val(1))
    cat32(x...)=cat(x...; dims=Val(2))
    cat33(x...)=cat(x...; dims=Val(3))

    for T in (Float32, Float64)
        x1d = (randn(T), randn(T))
        x2d = (randn(T,2), randn(T,2))
        x3d = (randn(T,2,3,5),randn(T,1,3,5),randn(T,2,2,5),randn(T,2,3,2))

        @test gradcheck(cat1, x1d...; o...)
        @test gradcheck(cat1, x2d...; o...)
        @test gradcheck(cat1, x1d[1], x2d[1]; o...)

        @test gradcheck(cat2, x1d...; o...)
        @test gradcheck(cat2, x2d...; o...)
        @test gradcheck(cat2, x1d[1], x2d[1]'; o...)

        @test gradcheck(cat31, x3d[1],x3d[2]; o...)
        @test gradcheck(cat32, x3d[1],x3d[3]; o...)
        @test gradcheck(cat33, x3d[1],x3d[4]; o...)

        @test gradcheck(cat1d,T[1.,2.],T[3.,4.]; o...)
    end
end


text:
### A Pluto.jl notebook ###
# v0.16.4

using Markdown
using InteractiveUtils

# ╔═╡ 2b00f9f2-53f3-4525-a3e4-708109259b6c
# Julia environment -
begin
	using PlutoUI
	using PrettyTables
end

# ╔═╡ c7822c42-2f57-11ec-3b48-2b2ea80c6f23
md"""$(PlutoUI.LocalResource("./figs/Fig-ModRankine-Heat-Recycle-F20.png"))
__Fig 1.__ Schematic of the unit operations of a modified organic Rankine cycle in which α percent of the heat from the condenser is recycled to the boiler through a heat exchanger."""

# ╔═╡ 796a9adb-7573-4ec1-881f-ab39d98032ae
md"""
The Organic Rankine Cycle (ORC) is an _open_ four step thermodynamic process used to generate power that uses organic compounds as working fluids (Fig. 1). In the cycle, path $\mathcal{P}_{ij}$ connects operating point $\mathcal{O}_{i}$ to $\mathcal{O}_{j}$:

* Path $\mathcal{P}_{12}$: $\left(1~\rightarrow~2\right)$: _isobaric_ heating in a boiler from operating point $\mathcal{O}_{1}$ to $\mathcal{O}_{2}$
* Path $\mathcal{P}_{23}$: $\left(2~\rightarrow~3\right)$: _adiabatic_ expansion in a turbine from operating point $\mathcal{O}_{2}$ to $\mathcal{O}_{3}$
* Path $\mathcal{P}_{34}$: $\left(3~\rightarrow~4\right)$: _isobaric_ cooling in a condenser from operating point $\mathcal{O}_{3}$ to $\mathcal{O}_{4}$
* Path $\mathcal{P}_{41}$: $\left(4~\rightarrow~1\right)$: _adiabatic_ compression in a pump from operating point $\mathcal{O}_{4}$ to $\mathcal{O}_{1}$

"""

# ╔═╡ d65087c7-1b55-4aeb-bda6-fd4039b64863
md"""

__Assumptions:__

* (i) the cycle operates at steady-state;
* (ii) the working fluid R-508B (Suva 95) has a mass flow rate of $\dot{m}$ = 1.25 kg s$^{-1}$;
* (iii) the turbine efficiency is $\eta_{T}$ = 70\%;
* (iv) _neglect_ the enthalpy and temperature change from the pump (assume T$_{4}\simeq$~T$_{1}$ and H$_{4}\simeq$~H$_{1}$);
* (v) neglect changes in the kinetic and potential energy in the system and streams.

"""

# ╔═╡ 1d0dbac8-8a12-4ba0-862b-e7a2860e2b6e
md"""
###### Example State Table for the Mod Rankine Problem
You get the $H$, $S$, $T$ and $P$ values from the problem statement, and from the working fluid PH-diagram (or data sheets). When completing the state table, always assume _unless explicitly asked otherwise_ reversibility for turbine.
"""

# ╔═╡ f561e6bd-63d0-43db-9bcb-5135afdfe7d4
begin

# 	# Version: δ
# 	# setup the state table -
# 	state_table_data_array = zeros(4,5)

# 	# setup each row -
# 	# row 1:
# 	state_table_data_array[1,1] = -100.0
# 	state_table_data_array[1,2] = 2200.0
# 	state_table_data_array[1,3] = 78.1
# 	state_table_data_array[1,4] = 0.4716
# 	state_table_data_array[1,5] = 0.0

# 	# row 2:
# 	state_table_data_array[2,1] = 20.0
# 	state_table_data_array[2,2] = 2200.0
# 	state_table_data_array[2,3] = 307.2
# 	state_table_data_array[2,4] = 1.4038
# 	state_table_data_array[2,5] = 1.0

# 	# row 3:
# 	state_table_data_array[3,1] = -100.0
# 	state_table_data_array[3,2] = 48.0
# 	state_table_data_array[3,3] = 240.3
# 	state_table_data_array[3,4] = 1.4038
# 	state_table_data_array[3,5] = 0.97

# 	# row 4:
# 	state_table_data_array[4,1] = -100.0
# 	state_table_data_array[4,2] = 48.0
# 	state_table_data_array[4,3] = 78.1
# 	state_table_data_array[4,4] = 0.4716
# 	state_table_data_array[4,5] = 0.0

	# show -
	nothing
end

# ╔═╡ 2cef6853-6624-428a-81c4-4e8fdb8727f7
begin

	# Version: γ (don't forget to change mdot, ηₜ and α)

	# setup the state table -
	state_table_data_array = zeros(4,5)

	# setup each row -
	# row 1:
	state_table_data_array[1,1] = -90.0
	state_table_data_array[1,2] = 2200.0
	state_table_data_array[1,3] = 86.8
	state_table_data_array[1,4] = 0.5184
	state_table_data_array[1,5] = 0.0

	# row 2:
	state_table_data_array[2,1] = 20.0
	state_table_data_array[2,2] = 2200.0
	state_table_data_array[2,3] = 307.2
	state_table_data_array[2,4] = 1.4038
	state_table_data_array[2,5] = 1.0

	# row 3:
	state_table_data_array[3,1] = -90.0
	state_table_data_array[3,2] = 90.0
	state_table_data_array[3,3] = 249.
	state_table_data_array[3,4] = 1.4038
	state_table_data_array[3,5] = 0.99

	# row 4:
	state_table_data_array[4,1] = -90.0
	state_table_data_array[4,2] = 90.0
	state_table_data_array[4,3] = 86.8
	state_table_data_array[4,4] = 0.5184
	state_table_data_array[4,5] = 0.0

	# show -
	nothing
end

# ╔═╡ 50bcce4a-981b-4b70-9fa6-98ad9f7b39e4
with_terminal() do

	# header row -
	state_table_header_row = (["T","P","H","S","θ"],["°C","kPa","kJ/kg","kJ/kg-K","dimensionless"]);

	# write the table -
	pretty_table(state_table_data_array; header=state_table_header_row)
end

# ╔═╡ d2bd39f1-8d81-404d-bea6-70ea8e19cea1
# do this computation in matrix - vector form
begin

	# setup process connectivity array A - change for γ and δ
	m_dot = 1.5 	# units: kg/s
	ηₜ = 0.75 		# units: dimensionless
	A = [-m_dot m_dot 0; 0 -m_dot m_dot ; m_dot 0 -m_dot];

	# get the H's
	H₁ = state_table_data_array[1,3] # units: kJ/kg
	H₂ = state_table_data_array[2,3] # units: kJ/kg
	H₃ = state_table_data_array[3,3] # units: kJ/kg
	H₄ = state_table_data_array[4,3] # units: kJ/kg
	H_vec = [H₁ ; H₂ ; H₃];

	# compute the RHS -
	bV = A*H_vec
end

# ╔═╡ 37101508-f1a7-42ed-b446-a9b21001c7b0
with_terminal() do

	# path table -
	path_table_data_array = zeros(5,3)

	# values -
	# row 1:
	path_table_data_array[1,1] = bV[1]
	path_table_data_array[1,2] = 0.0
	path_table_data_array[1,3] = 0.0

	# row 2:
	path_table_data_array[2,1] = 0.0
	path_table_data_array[2,2] = bV[2]
	path_table_data_array[2,3] = (ηₜ)*bV[2]

	# row 3:
	path_table_data_array[3,1] = bV[3]
	path_table_data_array[3,2] = 0.0
	path_table_data_array[3,3] = 0.0

	# row 4:
	path_table_data_array[4,1] = 0.0
	path_table_data_array[4,2] = 0.0
	path_table_data_array[4,3] = 0.0

	# row 5:
	path_table_data_array[5,1] = bV[1]+bV[3]
	path_table_data_array[5,2] = bV[2]
	path_table_data_array[5,3] = 0.0

	# header row -
	path_table_header_row = (["Qdot","Wdot (ideal)","Wdot (actual)"],["kJ/s","kJ/s","kJ/s"]);

	# write the table -
	pretty_table(path_table_data_array; header=path_table_header_row)
end

# ╔═╡ eab84d4f-1a60-478e-9d3e-92d1b466e142
md"""
###### Efficiency w/no heat recycle
"""

# ╔═╡ c6eca5d9-48b9-48bf-9bfa-de7359549e00
eff_no_recyle = -1*(bV[2])/(bV[1])

# ╔═╡ b7153ddf-b153-4b9c-9fea-aab1fe30471d
md"""
###### Efficiency w/heat recycle
"""

# ╔═╡ ac0f2a28-a8aa-4031-b69c-71e8f6047574
begin

	# what is the fraction of heat that we are going to recycle?
	α = 0.10

	# What is the amount of recycled heat -
	Q_dot_recycled = α*abs(bV[3])

	# what is the new amount of heat put into the boiler?
	Q_dot_new = bV[1] - Q_dot_recycled

	# compute new efficiency -
	eff_recyle = -1*(bV[2])/(Q_dot_new)
end

# ╔═╡ 15dbf654-25cc-43e5-8526-d3e62ca3271a
html"""
<style>
main {
    max-width: 1200px;
    width: 95%;
    margin: auto;
    font-family: "Roboto, monospace";
}

a {
    color: blue;
    text-decoration: none;
}

.H1 {
    padding: 0px 30px;
}
</style>"""

# ╔═╡ 00000000-0000-0000-0000-000000000001
PLUTO_PROJECT_TOML_CONTENTS = """
[deps]
PlutoUI = "7f904dfe-b85e-4ff6-b463-dae2292396a8"
PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"

[compat]
PlutoUI = "~0.7.16"
PrettyTables = "~1.2.3"
"""

# ╔═╡ 00000000-0000-0000-0000-000000000002
PLUTO_MANIFEST_TOML_CONTENTS = """
# This file is machine-generated - editing it directly is not advised

[[Base64]]
uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"

[[Crayons]]
git-tree-sha1 = "3f71217b538d7aaee0b69ab47d9b7724ca8afa0d"
uuid = "a8cc5b0e-0ffa-5ad4-8c14-923d3ee1735f"
version = "4.0.4"

[[DataAPI]]
git-tree-sha1 = "cc70b17275652eb47bc9e5f81635981f13cea5c8"
uuid = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a"
version = "1.9.0"

[[DataValueInterfaces]]
git-tree-sha1 = "bfc1187b79289637fa0ef6d4436ebdfe6905cbd6"
uuid = "e2d170a0-9d28-54be-80f0-106bbe20a464"
version = "1.0.0"

[[Dates]]
deps = ["Printf"]
uuid = "ade2ca70-3891-5945-98fb-dc099432e06a"

[[Formatting]]
deps = ["Printf"]
git-tree-sha1 = "8339d61043228fdd3eb658d86c926cb282ae72a8"
uuid = "59287772-0a20-5a39-b81b-1366585eb4c0"
version = "0.4.2"

[[Hyperscript]]
deps = ["Test"]
git-tree-sha1 = "8d511d5b81240fc8e6802386302675bdf47737b9"
uuid = "47d2ed2b-36de-50cf-bf87-49c2cf4b8b91"
version = "0.0.4"

[[HypertextLiteral]]
git-tree-sha1 = "f6532909bf3d40b308a0f360b6a0e626c0e263a8"
uuid = "ac1192a8-f4b3-4bfe-ba22-af5b92cd3ab2"
version = "0.9.1"

[[IOCapture]]
deps = ["Logging", "Random"]
git-tree-sha1 = "f7be53659ab06ddc986428d3a9dcc95f6fa6705a"
uuid = "b5f81e59-6552-4d32-b1f0-c071b021bf89"
version = "0.2.2"

[[InteractiveUtils]]
deps = ["Markdown"]
uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240"

[[IteratorInterfaceExtensions]]
git-tree-sha1 = "a3f24677c21f5bbe9d2a714f95dcd58337fb2856"
uuid = "82899510-4779-5014-852e-03e436cf321d"
version = "1.0.0"

[[JSON]]
deps = ["Dates", "Mmap", "Parsers", "Unicode"]
git-tree-sha1 = "8076680b162ada2a031f707ac7b4953e30667a37"
uuid = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
version = "0.21.2"

[[Libdl]]
uuid = "8f399da3-3557-5675-b5ff-fb832c97cbdb"

[[LinearAlgebra]]
deps = ["Libdl"]
uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

[[Logging]]
uuid = "56ddb016-857b-54e1-b83d-db4d58db5568"

[[Markdown]]
deps = ["Base64"]
uuid = "d6f4376e-aef5-505a-96c1-9c027394607a"

[[Mmap]]
uuid = "a63ad114-7e13-5084-954f-fe012c677804"

[[Parsers]]
deps = ["Dates"]
git-tree-sha1 = "98f59ff3639b3d9485a03a72f3ab35bab9465720"
uuid = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0"
version = "2.0.6"

[[PlutoUI]]
deps = ["Base64", "Dates", "Hyperscript", "HypertextLiteral", "IOCapture", "InteractiveUtils", "JSON", "Logging", "Markdown", "Random", "Reexport", "UUIDs"]
git-tree-sha1 = "4c8a7d080daca18545c56f1cac28710c362478f3"
uuid = "7f904dfe-b85e-4ff6-b463-dae2292396a8"
version = "0.7.16"

[[PrettyTables]]
deps = ["Crayons", "Formatting", "Markdown", "Reexport", "Tables"]
git-tree-sha1 = "d940010be611ee9d67064fe559edbb305f8cc0eb"
uuid = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
version = "1.2.3"

[[Printf]]
deps = ["Unicode"]
uuid = "de0858da-6303-5e67-8744-51eddeeeb8d7"

[[Random]]
deps = ["Serialization"]
uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

[[Reexport]]
git-tree-sha1 = "45e428421666073eab6f2da5c9d310d99bb12f9b"
uuid = "189a3867-3050-52da-a836-e630ba90ab69"
version = "1.2.2"

[[SHA]]
uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce"

[[Serialization]]
uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b"

[[TableTraits]]
deps = ["IteratorInterfaceExtensions"]
git-tree-sha1 = "c06b2f539df1c6efa794486abfb6ed2022561a39"
uuid = "3783bdb8-4a98-5b6b-af9a-565f29a5fe9c"
version = "1.0.1"

[[Tables]]
deps = ["DataAPI", "DataValueInterfaces", "IteratorInterfaceExtensions", "LinearAlgebra", "TableTraits", "Test"]
git-tree-sha1 = "fed34d0e71b91734bf0a7e10eb1bb05296ddbcd0"
uuid = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
version = "1.6.0"

[[Test]]
deps = ["InteractiveUtils", "Logging", "Random", "Serialization"]
uuid = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

[[UUIDs]]
deps = ["Random", "SHA"]
uuid = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"

[[Unicode]]
uuid = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"
"""

# ╔═╡ Cell order:
# ╟─c7822c42-2f57-11ec-3b48-2b2ea80c6f23
# ╟─796a9adb-7573-4ec1-881f-ab39d98032ae
# ╟─d65087c7-1b55-4aeb-bda6-fd4039b64863
# ╟─1d0dbac8-8a12-4ba0-862b-e7a2860e2b6e
# ╠═f561e6bd-63d0-43db-9bcb-5135afdfe7d4
# ╠═2cef6853-6624-428a-81c4-4e8fdb8727f7
# ╠═50bcce4a-981b-4b70-9fa6-98ad9f7b39e4
# ╠═d2bd39f1-8d81-404d-bea6-70ea8e19cea1
# ╠═37101508-f1a7-42ed-b446-a9b21001c7b0
# ╟─eab84d4f-1a60-478e-9d3e-92d1b466e142
# ╠═c6eca5d9-48b9-48bf-9bfa-de7359549e00
# ╟─b7153ddf-b153-4b9c-9fea-aab1fe30471d
# ╠═ac0f2a28-a8aa-4031-b69c-71e8f6047574
# ╠═2b00f9f2-53f3-4525-a3e4-708109259b6c
# ╟─15dbf654-25cc-43e5-8526-d3e62ca3271a
# ╟─00000000-0000-0000-0000-000000000001
# ╟─00000000-0000-0000-0000-000000000002


text:
using Printf

struct KPCSupply{T<:Real}
    item::String
    weight::T
    value::T
    uvalue::T
end
function KPCSupply(item::AbstractString, weight::Real, value::Real)
    w, v = promote(weight, value)
    KPCSupply(item, w, v, v / w)
end

Base.show(io::IO, s::KPCSupply) = print(io, s.item, @sprintf " (%.2f kg, %.2f €, %.2f €/kg)" s.weight s.value s.uvalue)
Base.isless(a::KPCSupply, b::KPCSupply) = a.uvalue < b.uvalue

function solve(store::Vector{KPCSupply{T}}, capacity::Real) where T<:Real
    sack = similar(store, 0) # vector like store, but of length 0
    kweight = zero(T)
    for s in sort(store, rev = true)
        if kweight + s.weight ≤ capacity
            kweight += s.weight
            push!(sack, s)
        else
            w = capacity - kweight
            v = w * s.uvalue
            push!(sack, KPCSupply(s.item, w, v, s.value))
            break
        end
    end
    return sack
end

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
function missingperm1(arr::Vector{<:AbstractString})
    missperm = string()
    for pos in 1:length(arr[1])
        s = Set()
        for perm in arr
            c = perm[pos]
            if c ∈ s pop!(s, c) else push!(s, c) end
        end
        missperm *= first(s)
    end
    return missperm
end


text:
# SPDX-License-Identifier: MIT

module MarkovModels

using CUDA
using CUDA.CUSPARSE
using LinearAlgebra
using SparseArrays

#======================================================================
Redefinition of some linear operations to work on GPU.
======================================================================#

include("linalg.jl")
include("culinalg.jl")

#======================================================================
Semifield algebras.
======================================================================#

include("semifields.jl")

export LogSemifield
export TropicalSemifield

#======================================================================
API to build and manipulate FSM.
======================================================================#

include("fsm.jl")

export FSM
export addstate!
export determinize
export link!
export minimize
export renormalize!
export setinit!
export setfinal!
export states
export remove_eps
export fsmtype

include("cfsm.jl")

export CompiledFSM
export compile
export gpu
export fsmtype

#======================================================================
Inference algorithms.
======================================================================#

include("algorithms.jl")

export stateposteriors
export bestpath

end

max_stars_repo_name: MartinKocour/MarkovModels

max_issues_repo_name: MartinKocour/MarkovModels

max_forks_repo_name: MartinKocour/MarkovModels


text:
#### EDMF model kernels

using CLIMAParameters.Planet: e_int_v0, grav, day, R_d, R_v, molmass_ratio
using Printf
using ClimateMachine.Atmos: nodal_update_auxiliary_state!, Advect

using ClimateMachine.BalanceLaws

using ClimateMachine.MPIStateArrays: MPIStateArray
using ClimateMachine.DGMethods: LocalGeometry, DGModel

import ClimateMachine.Atmos: atmos_source!

import ClimateMachine.BalanceLaws:
    vars_state,
    prognostic_vars,
    flux,
    source,
    eq_tends,
    update_auxiliary_state!,
    init_state_prognostic!,
    flux_first_order!,
    flux_second_order!,
    compute_gradient_argument!,
    compute_gradient_flux!

import ClimateMachine.TurbulenceConvection:
    init_aux_turbconv!,
    turbconv_nodal_update_auxiliary_state!,
    turbconv_boundary_state!,
    turbconv_normal_boundary_flux_second_order!

using ClimateMachine.Thermodynamics: air_pressure, air_density

include(joinpath("helper_funcs", "nondimensional_exchange_functions.jl"))
include(joinpath("helper_funcs", "lamb_smooth_minimum.jl"))
include(joinpath("helper_funcs", "utility_funcs.jl"))
include(joinpath("helper_funcs", "subdomain_statistics.jl"))
include(joinpath("helper_funcs", "diagnose_environment.jl"))
include(joinpath("helper_funcs", "subdomain_thermo_states.jl"))
include(joinpath("helper_funcs", "save_subdomain_temperature.jl"))
include(joinpath("closures", "entr_detr.jl"))
include(joinpath("closures", "pressure.jl"))
include(joinpath("closures", "mixing_length.jl"))
include(joinpath("closures", "turbulence_functions.jl"))
include(joinpath("closures", "surface_functions.jl"))

function vars_state(m::NTuple{N, Updraft}, st::Auxiliary, FT) where {N}
    return Tuple{ntuple(i -> vars_state(m[i], st, FT), N)...}
end

function vars_state(::Updraft, ::Auxiliary, FT)
    @vars(
        buoyancy::FT,
        a::FT,
        E_dyn::FT,
        Δ_dyn::FT,
        E_trb::FT,
        T::FT,
        θ_liq::FT,
        q_tot::FT,
        w::FT,
    )
end

function vars_state(::Environment, ::Auxiliary, FT)
    @vars(T::FT, cld_frac::FT, buoyancy::FT)
end

function vars_state(m::EDMF, st::Auxiliary, FT)
    @vars(
        environment::vars_state(m.environment, st, FT),
        updraft::vars_state(m.updraft, st, FT)
    )
end

function vars_state(::Updraft, ::Prognostic, FT)
    @vars(ρa::FT, ρaw::FT, ρaθ_liq::FT, ρaq_tot::FT,)
end

function vars_state(::Environment, ::Prognostic, FT)
    @vars(ρatke::FT, ρaθ_liq_cv::FT, ρaq_tot_cv::FT, ρaθ_liq_q_tot_cv::FT,)
end

function vars_state(m::NTuple{N, Updraft}, st::Prognostic, FT) where {N}
    return Tuple{ntuple(i -> vars_state(m[i], st, FT), N)...}
end

function vars_state(m::EDMF, st::Prognostic, FT)
    @vars(
        environment::vars_state(m.environment, st, FT),
        updraft::vars_state(m.updraft, st, FT)
    )
end

function vars_state(::Updraft, ::Gradient, FT)
    @vars(w::FT,)
end

function vars_state(::Environment, ::Gradient, FT)
    @vars(
        θ_liq::FT,
        q_tot::FT,
        w::FT,
        tke::FT,
        θ_liq_cv::FT,
        q_tot_cv::FT,
        θ_liq_q_tot_cv::FT,
        θv::FT,
        e::FT,
    )
end

function vars_state(m::NTuple{N, Updraft}, st::Gradient, FT) where {N}
    return Tuple{ntuple(i -> vars_state(m[i], st, FT), N)...}
end

function vars_state(m::EDMF, st::Gradient, FT)
    @vars(
        environment::vars_state(m.environment, st, FT),
        updraft::vars_state(m.updraft, st, FT)
    )
end

function vars_state(m::NTuple{N, Updraft}, st::GradientFlux, FT) where {N}
    return Tuple{ntuple(i -> vars_state(m[i], st, FT), N)...}
end

function vars_state(::Updraft, st::GradientFlux, FT)
    @vars(∇w::SVector{3, FT},)
end

function vars_state(::Environment, ::GradientFlux, FT)
    @vars(
        ∇θ_liq::SVector{3, FT},
        ∇q_tot::SVector{3, FT},
        ∇w::SVector{3, FT},
        ∇tke::SVector{3, FT},
        ∇θ_liq_cv::SVector{3, FT},
        ∇q_tot_cv::SVector{3, FT},
        ∇θ_liq_q_tot_cv::SVector{3, FT},
        ∇θv::SVector{3, FT},
        ∇e::SVector{3, FT},
    )

end

function vars_state(m::EDMF, st::GradientFlux, FT)
    @vars(
        S²::FT, # should be conditionally grabbed from atmos.turbulence
        environment::vars_state(m.environment, st, FT),
        updraft::vars_state(m.updraft, st, FT)
    )
end

abstract type EDMFPrognosticVariable <: PrognosticVariable end

abstract type EnvironmentPrognosticVariable <: EDMFPrognosticVariable end
struct en_ρatke <: EnvironmentPrognosticVariable end
struct en_ρaθ_liq_cv <: EnvironmentPrognosticVariable end
struct en_ρaq_tot_cv <: EnvironmentPrognosticVariable end
struct en_ρaθ_liq_q_tot_cv <: EnvironmentPrognosticVariable end

abstract type UpdraftPrognosticVariable{i} <: EDMFPrognosticVariable end
struct up_ρa{i} <: UpdraftPrognosticVariable{i} end
struct up_ρaw{i} <: UpdraftPrognosticVariable{i} end
struct up_ρaθ_liq{i} <: UpdraftPrognosticVariable{i} end
struct up_ρaq_tot{i} <: UpdraftPrognosticVariable{i} end

prognostic_vars(m::EDMF) =
    (prognostic_vars(m.environment)..., prognostic_vars(m.updraft)...)
prognostic_vars(m::Environment) =
    (en_ρatke(), en_ρaθ_liq_cv(), en_ρaq_tot_cv(), en_ρaθ_liq_q_tot_cv())

function prognostic_vars(m::NTuple{N, Updraft}) where {N}
    t_ρa = ntuple(i -> up_ρa{i}(), N)
    t_ρaw = ntuple(i -> up_ρaw{i}(), N)
    t_ρaθ_liq = ntuple(i -> up_ρaθ_liq{i}(), N)
    t_ρaq_tot = ntuple(i -> up_ρaq_tot{i}(), N)
    t = (t_ρa..., t_ρaw..., t_ρaθ_liq..., t_ρaq_tot...)
    return t
end

# Dycore tendencies
eq_tends(
    pv::PV,
    m::EDMF,
    ::Flux{SecondOrder},
) where {PV <: Union{Momentum, Energy, TotalMoisture}} = ()
# (SGSFlux{PV}(),) # to add SGSFlux back to grid-mean

# Turbconv tendencies
eq_tends(
    pv::PV,
    m::AtmosModel,
    tt::Flux{O},
) where {O, PV <: EDMFPrognosticVariable} = eq_tends(pv, m.turbconv, tt)

eq_tends(pv::PV, m::EDMF, ::Flux{O}) where {O, PV <: EDMFPrognosticVariable} =
    ()

eq_tends(
    pv::PV,
    m::EDMF,
    ::Flux{FirstOrder},
) where {PV <: EDMFPrognosticVariable} = (Advect{PV}(),)

struct SGSFlux{PV <: Union{Momentum, Energy, TotalMoisture}} <:
       TendencyDef{Flux{SecondOrder}, PV} end

"""
    init_aux_turbconv!(
        turbconv::EDMF{FT},
        m::AtmosModel{FT},
        aux::Vars,
        geom::LocalGeometry,
    ) where {FT}

Initialize EDMF auxiliary variables.
"""
function init_aux_turbconv!(
    turbconv::EDMF{FT},
    m::AtmosModel{FT},
    aux::Vars,
    geom::LocalGeometry,
) where {FT}
    N_up = n_updrafts(turbconv)

    # Aliases:
    en_aux = aux.turbconv.environment
    up_aux = aux.turbconv.updraft

    en_aux.cld_frac = FT(0)
    en_aux.buoyancy = FT(0)

    @unroll_map(N_up) do i
        up_aux[i].buoyancy = FT(0)
        up_aux[i].θ_liq = FT(0)
        up_aux[i].q_tot = FT(0)
        up_aux[i].w = FT(0)
    end
end;

function turbconv_nodal_update_auxiliary_state!(
    turbconv::EDMF{FT},
    m::AtmosModel{FT},
    state::Vars,
    aux::Vars,
    t::Real,
) where {FT}
    N_up = n_updrafts(turbconv)
    save_subdomain_temperature!(m, state, aux)

    en_aux = aux.turbconv.environment
    up_aux = aux.turbconv.updraft
    gm = state
    en = state.turbconv.environment
    up = state.turbconv.updraft

    # Recover thermo states
    ts = recover_thermo_state_all(m, state, aux)

    # Get environment variables
    env = environment_vars(state, aux, N_up)

    # Compute buoyancies of subdomains
    ρ_inv = 1 / gm.ρ
    _grav::FT = grav(m.param_set)

    z = altitude(m, aux)

    ρ_en = air_density(ts.en)
    en_aux.buoyancy = -_grav * (ρ_en - aux.ref_state.ρ) * ρ_inv

    @unroll_map(N_up) do i
        ρ_i = air_density(ts.up[i])
        up_aux[i].buoyancy = -_grav * (ρ_i - aux.ref_state.ρ) * ρ_inv
        up_aux[i].a = up[i].ρa * ρ_inv
        up_aux[i].θ_liq = up[i].ρaθ_liq / up[i].ρa
        up_aux[i].q_tot = up[i].ρaq_tot / up[i].ρa
        up_aux[i].w = up[i].ρaw / up[i].ρa
    end
    b_gm = grid_mean_b(state, aux, N_up)

    # remove the gm_b from all subdomains
    @unroll_map(N_up) do i
        up_aux[i].buoyancy -= b_gm
    end
    en_aux.buoyancy -= b_gm

    EΔ_up = ntuple(N_up) do i
        entr_detr(m, m.turbconv.entr_detr, state, aux, t, ts, env, i)
    end

    E_dyn, Δ_dyn, E_trb = ntuple(i -> map(x -> x[i], EΔ_up), 3)

    @unroll_map(N_up) do i
        up_aux[i].E_dyn = E_dyn[i]
        up_aux[i].Δ_dyn = Δ_dyn[i]
        up_aux[i].E_trb = E_trb[i]
    end

end;

function compute_gradient_argument!(
    turbconv::EDMF{FT},
    m::AtmosModel{FT},
    transform::Vars,
    state::Vars,
    aux::Vars,
    t::Real,
) where {FT}
    N_up = n_updrafts(turbconv)
    z = altitude(m, aux)

    # Aliases:
    up_tf = transform.turbconv.updraft
    en_tf = transform.turbconv.environment
    gm = state
    up = state.turbconv.updraft
    en = state.turbconv.environment

    # Recover thermo states
    ts = recover_thermo_state_all(m, state, aux)

    # Get environment variables
    env = environment_vars(state, aux, N_up)

    @unroll_map(N_up) do i
        up_tf[i].w = up[i].ρaw / up[i].ρa
    end
    _grav::FT = grav(m.param_set)

    ρ_inv = 1 / gm.ρ
    θ_liq_en = liquid_ice_pottemp(ts.en)
    q_tot_en = total_specific_humidity(ts.en)

    # populate gradient arguments
    en_tf.θ_liq = θ_liq_en
    en_tf.q_tot = q_tot_en
    en_tf.w = env.w

    en_tf.tke = enforce_positivity(en.ρatke) / (env.a * gm.ρ)
    en_tf.θ_liq_cv = enforce_positivity(en.ρaθ_liq_cv) / (env.a * gm.ρ)
    en_tf.q_tot_cv = enforce_positivity(en.ρaq_tot_cv) / (env.a * gm.ρ)
    en_tf.θ_liq_q_tot_cv = en.ρaθ_liq_q_tot_cv / (env.a * gm.ρ)

    # TODO: is this supposed to be grabbed from grid mean?
    en_tf.θv = virtual_pottemp(ts.gm)
    e_kin = FT(1 // 2) * ((gm.ρu[1] * ρ_inv)^2 + (gm.ρu[2] * ρ_inv)^2 + env.w^2)
    en_tf.e = total_energy(e_kin, _grav * z, ts.en)
end;

function compute_gradient_flux!(
    turbconv::EDMF{FT},
    m::AtmosModel{FT},
    diffusive::Vars,
    ∇transform::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
) where {FT}
    N_up = n_updrafts(turbconv)

    # Aliases:
    gm = state
    up_dif = diffusive.turbconv.updraft
    tc_dif = diffusive.turbconv
    up_∇tf = ∇transform.turbconv.updraft
    en_dif = diffusive.turbconv.environment
    en_∇tf = ∇transform.turbconv.environment

    @unroll_map(N_up) do i
        up_dif[i].∇w = up_∇tf[i].w
    end

    ρ_inv = 1 / gm.ρ
    # first moment grid mean coming from environment gradients only
    en_dif.∇θ_liq = en_∇tf.θ_liq
    en_dif.∇q_tot = en_∇tf.q_tot
    en_dif.∇w = en_∇tf.w
    # second moment env cov
    en_dif.∇tke = en_∇tf.tke
    en_dif.∇θ_liq_cv = en_∇tf.θ_liq_cv
    en_dif.∇q_tot_cv = en_∇tf.q_tot_cv
    en_dif.∇θ_liq_q_tot_cv = en_∇tf.θ_liq_q_tot_cv

    en_dif.∇θv = en_∇tf.θv
    en_dif.∇e = en_∇tf.e

    tc_dif.S² = ∇transform.u[3, 1]^2 + ∇transform.u[3, 2]^2 + en_dif.∇w[3]^2
end;

struct TurbconvSource <: AbstractSource end

function atmos_source!(
    ::TurbconvSource,
    m::AtmosModel{FT},
    source::Vars,
    state::Vars,
    diffusive::Vars,
    aux::Vars,
    t::Real,
    direction,
) where {FT}
    turbconv = m.turbconv
    N_up = n_updrafts(turbconv)

    # Aliases:
    gm = state
    en = state.turbconv.environment
    up = state.turbconv.updraft
    en_src = source.turbconv.environment
    up_src = source.turbconv.updraft
    en_dif = diffusive.turbconv.environment
    up_aux = aux.turbconv.updraft

    # Recover thermo states
    ts = recover_thermo_state_all(m, state, aux)

    # Get environment variables
    env = environment_vars(state, aux, N_up)

    EΔ_up = ntuple(N_up) do i
        entr_detr(m, m.turbconv.entr_detr, state, aux, t, ts, env, i)
    end
    E_dyn, Δ_dyn, E_trb = ntuple(i -> map(x -> x[i], EΔ_up), 3)

    # get environment values
    _grav::FT = grav(m.param_set)
    ρ_inv = 1 / gm.ρ
    θ_liq_en = liquid_ice_pottemp(ts.en)
    q_tot_en = total_specific_humidity(ts.en)
    tke_en = enforce_positivity(en.ρatke) * ρ_inv / env.a
    θ_liq = liquid_ice_pottemp(ts.gm)
    a_min = turbconv.subdomains.a_min
    a_max = turbconv.subdomains.a_max

    ρa_up = vuntuple(N_up) do i
        gm.ρ * enforce_unit_bounds(up[i].ρa * ρ_inv, a_min, a_max)
    end
    ρq_tot = m.moisture isa DryModel ? FT(0) : gm.moisture.ρq_tot

    @unroll_map(N_up) do i

        ρa_up_i = ρa_up[i]
        w_up_i = up[i].ρaw / ρa_up_i
        ρa_up_i_inv = FT(1) / ρa_up_i

        # first moment sources - for now we compute these as aux variable
        dpdz = perturbation_pressure(
            m,
            m.turbconv.pressure,
            state,
            diffusive,
            aux,
            t,
            env,
            i,
        )

        # entrainment and detrainment
        up_src[i].ρa += E_dyn[i] - Δ_dyn[i]
        up_src[i].ρaw +=
            ((E_dyn[i] + E_trb[i]) * env.w - (Δ_dyn[i] + E_trb[i]) * w_up_i)
        up_src[i].ρaθ_liq += (
            (E_dyn[i] + E_trb[i]) * θ_liq_en -
            (Δ_dyn[i] + E_trb[i]) * up[i].ρaθ_liq * ρa_up_i_inv
        )
        up_src[i].ρaq_tot += (
            (E_dyn[i] + E_trb[i]) * q_tot_en -
            (Δ_dyn[i] + E_trb[i]) * up[i].ρaq_tot * ρa_up_i_inv
        )

        # add buoyancy and perturbation pressure in subdomain w equation
        up_src[i].ρaw += up[i].ρa * (up_aux[i].buoyancy - dpdz)
        # microphysics sources should be applied here

        # environment second moments:
        en_src.ρatke += (
            Δ_dyn[i] * (w_up_i - env.w) * (w_up_i - env.w) * FT(0.5) +
            E_trb[i] * (env.w - gm.ρu[3] * ρ_inv) * (env.w - w_up_i) -
            (E_dyn[i] + E_trb[i]) * tke_en
        )

        en_src.ρaθ_liq_cv += (
            Δ_dyn[i] *
            (up[i].ρaθ_liq * ρa_up_i_inv - θ_liq_en) *
            (up[i].ρaθ_liq * ρa_up_i_inv - θ_liq_en) +
            E_trb[i] *
            (θ_liq_en - θ_liq) *
            (θ_liq_en - up[i].ρaθ_liq * ρa_up_i_inv) +
            E_trb[i] *
            (θ_liq_en - θ_liq) *
            (θ_liq_en - up[i].ρaθ_liq * ρa_up_i_inv) -
            (E_dyn[i] + E_trb[i]) * en.ρaθ_liq_cv
        )

        en_src.ρaq_tot_cv += (
            Δ_dyn[i] *
            (up[i].ρaq_tot * ρa_up_i_inv - q_tot_en) *
            (up[i].ρaq_tot * ρa_up_i_inv - q_tot_en) +
            E_trb[i] *
            (q_tot_en - ρq_tot * ρ_inv) *
            (q_tot_en - up[i].ρaq_tot * ρa_up_i_inv) +
            E_trb[i] *
            (q_tot_en - ρq_tot * ρ_inv) *
            (q_tot_en - up[i].ρaq_tot * ρa_up_i_inv) -
            (E_dyn[i] + E_trb[i]) * en.ρaq_tot_cv
        )

        en_src.ρaθ_liq_q_tot_cv += (
            Δ_dyn[i] *
            (up[i].ρaθ_liq * ρa_up_i_inv - θ_liq_en) *
            (up[i].ρaq_tot * ρa_up_i_inv - q_tot_en) +
            E_trb[i] *
            (θ_liq_en - θ_liq) *
            (q_tot_en - up[i].ρaq_tot * ρa_up_i_inv) +
            E_trb[i] *
            (q_tot_en - ρq_tot * ρ_inv) *
            (θ_liq_en - up[i].ρaθ_liq * ρa_up_i_inv) -
            (E_dyn[i] + E_trb[i]) * en.ρaθ_liq_q_tot_cv
        )

        # pressure tke source from the i'th updraft
        en_src.ρatke += ρa_up_i * (w_up_i - env.w) * dpdz
    end
    l_mix, ∂b∂z_env, Pr_t = mixing_length(
        m,
        m.turbconv.mix_len,
        state,
        diffusive,
        aux,
        t,
        Δ_dyn,
        E_trb,
        ts,
        env,
    )

    K_m = m.turbconv.mix_len.c_m * l_mix * sqrt(tke_en)
    K_h = K_m / Pr_t
    Shear² = diffusive.turbconv.S²
    ρa₀ = gm.ρ * env.a
    Diss₀ = m.turbconv.mix_len.c_d * sqrt(tke_en) / l_mix

    # production from mean gradient and Dissipation
    en_src.ρatke += ρa₀ * K_m * Shear² # tke Shear source
    en_src.ρatke += -ρa₀ * K_h * ∂b∂z_env   # tke Bouyancy source
    en_src.ρatke += -ρa₀ * Diss₀ * tke_en  # tke Dissipation

    en_src.ρaθ_liq_cv +=
        ρa₀ *
        (K_h * en_dif.∇θ_liq[3] * en_dif.∇θ_liq[3] - Diss₀ * en.ρaθ_liq_cv)
    en_src.ρaq_tot_cv +=
        ρa₀ *
        (K_h * en_dif.∇q_tot[3] * en_dif.∇q_tot[3] - Diss₀ * en.ρaq_tot_cv)
    en_src.ρaθ_liq_q_tot_cv +=
        ρa₀ * (
            K_h * en_dif.∇θ_liq[3] * en_dif.∇q_tot[3] -
            Diss₀ * en.ρaθ_liq_q_tot_cv
        )
    # covariance microphysics sources should be applied here
end;

function compute_ρa_up(m, state, aux)
    # Aliases:
    turbconv = m.turbconv
    gm = state
    up = state.turbconv.updraft
    N_up = n_updrafts(turbconv)
    a_min = turbconv.subdomains.a_min
    a_max = turbconv.subdomains.a_max
    # in future GCM implementations we need to think about grid mean advection
    ρa_up = vuntuple(N_up) do i
        gm.ρ * enforce_unit_bounds(up[i].ρa / gm.ρ, a_min, a_max)
    end
    return ρa_up
end

function flux(::Advect{up_ρa{i}}, m, state, aux, t, ts, direction) where {i}
    up = state.turbconv.updraft
    ẑ = vertical_unit_vector(m, aux)
    return up[i].ρaw * ẑ
end
function flux(::Advect{up_ρaw{i}}, m, state, aux, t, ts, direction) where {i}
    up = state.turbconv.updraft
    ẑ = vertical_unit_vector(m, aux)
    ρa_up = compute_ρa_up(m, state, aux)
    return up[i].ρaw * up[i].ρaw / ρa_up[i] * ẑ
end
function flux(
    ::Advect{up_ρaθ_liq{i}},
    m,
    state,
    aux,
    t,
    ts,
    direction,
) where {i}
    up = state.turbconv.updraft
    ẑ = vertical_unit_vector(m, aux)
    ρa_up = compute_ρa_up(m, state, aux)
    return up[i].ρaw / ρa_up[i] * up[i].ρaθ_liq * ẑ
end
function flux(
    ::Advect{up_ρaq_tot{i}},
    m,
    state,
    aux,
    t,
    ts,
    direction,
) where {i}
    up = state.turbconv.updraft
    ẑ = vertical_unit_vector(m, aux)
    ρa_up = compute_ρa_up(m, state, aux)
    return up[i].ρaw / ρa_up[i] * up[i].ρaq_tot * ẑ
end

function flux(::Advect{en_ρatke}, m, state, aux, t, ts, direction)
    en = state.turbconv.environment
    env = environment_vars(state, aux, n_updrafts(m.turbconv))
    ẑ = vertical_unit_vector(m, aux)
    return en.ρatke * env.w * ẑ
end
function flux(::Advect{en_ρaθ_liq_cv}, m, state, aux, t, ts, direction)
    en = state.turbconv.environment
    env = environment_vars(state, aux, n_updrafts(m.turbconv))
    ẑ = vertical_unit_vector(m, aux)
    return en.ρaθ_liq_cv * env.w * ẑ
end
function flux(::Advect{en_ρaq_tot_cv}, m, state, aux, t, ts, direction)
    en = state.turbconv.environment
    env = environment_vars(state, aux, n_updrafts(m.turbconv))
    ẑ = vertical_unit_vector(m, aux)
    return en.ρaq_tot_cv * env.w * ẑ
end
function flux(::Advect{en_ρaθ_liq_q_tot_cv}, m, state, aux, t, ts, direction)
    en = state.turbconv.environment
    env = environment_vars(state, aux, n_updrafts(m.turbconv))
    ẑ = vertical_unit_vector(m, aux)
    return en.ρaθ_liq_q_tot_cv * env.w * ẑ
end

# # in the EDMF first order (advective) fluxes exist only in the grid mean (if <w> is nonzero) and the uprdafts
function flux_first_order!(
    turbconv::EDMF{FT},
    m::AtmosModel{FT},
    flux::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
    ts,
    direction,
) where {FT}
    # Aliases:
    up_flx = flux.turbconv.updraft
    en_flx = flux.turbconv.environment
    N_up = n_updrafts(turbconv)
    # in future GCM implementations we need to think about grid mean advection
    tend = Flux{FirstOrder}()
    args = (m, state, aux, t, ts, direction)

    @unroll_map(N_up) do i
        up_flx[i].ρa = Σfluxes(eq_tends(up_ρa{i}(), m, tend), args...)
        up_flx[i].ρaw = Σfluxes(eq_tends(up_ρaw{i}(), m, tend), args...)
        up_flx[i].ρaθ_liq = Σfluxes(eq_tends(up_ρaθ_liq{i}(), m, tend), args...)
        up_flx[i].ρaq_tot = Σfluxes(eq_tends(up_ρaq_tot{i}(), m, tend), args...)
    end
    en_flx.ρatke = Σfluxes(eq_tends(en_ρatke(), m, tend), args...)
    en_flx.ρaθ_liq_cv = Σfluxes(eq_tends(en_ρaθ_liq_cv(), m, tend), args...)
    en_flx.ρaq_tot_cv = Σfluxes(eq_tends(en_ρaq_tot_cv(), m, tend), args...)
    en_flx.ρaθ_liq_q_tot_cv =
        Σfluxes(eq_tends(en_ρaθ_liq_q_tot_cv(), m, tend), args...)
end;

# in the EDMF second order (diffusive) fluxes
# exist only in the grid mean and the environment
function flux_second_order!(
    turbconv::EDMF{FT},
    m::AtmosModel{FT},
    flux::Grad,
    state::Vars,
    diffusive::Vars,
    aux::Vars,
    t::Real,
) where {FT}
    N_up = n_updrafts(turbconv)

    # Aliases:
    gm = state
    up = state.turbconv.updraft
    en = state.turbconv.environment
    gm_flx = flux
    en_flx = flux.turbconv.environment
    en_dif = diffusive.turbconv.environment

    # Recover thermo states
    ts = recover_thermo_state_all(m, state, aux)

    # Get environment variables
    env = environment_vars(state, aux, N_up)

    ρ_inv = FT(1) / gm.ρ
    _grav::FT = grav(m.param_set)
    z = altitude(m, aux)
    a_min = turbconv.subdomains.a_min
    a_max = turbconv.subdomains.a_max

    EΔ_up = ntuple(N_up) do i
        entr_detr(m, m.turbconv.entr_detr, state, aux, t, ts, env, i)
    end

    E_dyn, Δ_dyn, E_trb = ntuple(i -> map(x -> x[i], EΔ_up), 3)

    l_mix, _, Pr_t = mixing_length(
        m,
        turbconv.mix_len,
        state,
        diffusive,
        aux,
        t,
        Δ_dyn,
        E_trb,
        ts,
        env,
    )
    tke_en = enforce_positivity(en.ρatke) / env.a * ρ_inv
    K_m = m.turbconv.mix_len.c_m * l_mix * sqrt(tke_en)
    K_h = K_m / Pr_t

    #TotalFlux(ϕ) = Eddy_Diffusivity(ϕ) + MassFlux(ϕ)
    e_int = internal_energy(m, state, aux)

    e_kin = vuntuple(N_up) do i
        FT(1 // 2) * (
            (gm.ρu[1] * ρ_inv)^2 +
            (gm.ρu[2] * ρ_inv)^2 +
            (up[i].ρaw / up[i].ρa)^2
        )
    end
    e_tot_up = ntuple(i -> total_energy(e_kin[i], _grav * z, ts.up[i]), N_up)
    ρa_up = vuntuple(N_up) do i
        gm.ρ * enforce_unit_bounds(up[i].ρa * ρ_inv, a_min, a_max)
    end

    massflux_e = sum(
        vuntuple(N_up) do i
            up[i].ρa *
            ρ_inv *
            (gm.ρe * ρ_inv - e_tot_up[i]) *
            (gm.ρu[3] * ρ_inv - up[i].ρaw / ρa_up[i])
        end,
    )
    ρq_tot = m.moisture isa DryModel ? FT(0) : gm.moisture.ρq_tot
    massflux_q_tot = sum(
        vuntuple(N_up) do i
            up[i].ρa *
            ρ_inv *
            (ρq_tot * ρ_inv - up[i].ρaq_tot / up[i].ρa) *
            (gm.ρu[3] * ρ_inv - up[i].ρaw / ρa_up[i])
        end,
    )

    massflux_w = sum(
        vuntuple(N_up) do i
            up[i].ρa *
            ρ_inv *
            (gm.ρu[3] * ρ_inv - up[i].ρaw / up[i].ρa) *
            (gm.ρu[3] * ρ_inv - up[i].ρaw / ρa_up[i])
        end,
    )

    # update grid mean flux_second_order
    ρe_sgs_flux = -gm.ρ * env.a * K_h * en_dif.∇e[3] + massflux_e
    ρq_tot_sgs_flux = -gm.ρ * env.a * K_h * en_dif.∇q_tot[3] + massflux_q_tot
    ρu_sgs_flux = -gm.ρ * env.a * K_m * en_dif.∇w[3] + massflux_w

    # for now the coupling to the dycore is commented out

    # gm_flx.ρe              += SVector{3,FT}(0,0,ρe_sgs_flux)
    # gm_flx.moisture.ρq_tot += SVector{3,FT}(0,0,ρq_tot_sgs_flux)
    # gm_flx.ρu              += SMatrix{3, 3, FT, 9}(
    #     0, 0, 0,
    #     0, 0, 0,
    #     0, 0, ρu_sgs_flux,
    # )

    ẑ = vertical_unit_vector(m, aux)
    # env second moment flux_second_order
    en_flx.ρatke = -gm.ρ * env.a * K_m * en_dif.∇tke[3] * ẑ
    en_flx.ρaθ_liq_cv = -gm.ρ * env.a * K_h * en_dif.∇θ_liq_cv[3] * ẑ
    en_flx.ρaq_tot_cv = -gm.ρ * env.a * K_h * en_dif.∇q_tot_cv[3] * ẑ
    en_flx.ρaθ_liq_q_tot_cv =
        -gm.ρ * env.a * K_h * en_dif.∇θ_liq_q_tot_cv[3] * ẑ
end;

# First order boundary conditions
function turbconv_boundary_state!(
    nf,
    bc::EDMFBottomBC,
    m::AtmosModel{FT},
    state⁺::Vars,
    aux⁺::Vars,
    n⁻,
    state⁻::Vars,
    aux⁻::Vars,
    t,
    state_int::Vars,
    aux_int::Vars,
) where {FT}

    turbconv = m.turbconv
    N_up = n_updrafts(turbconv)
    up = state⁺.turbconv.updraft
    en = state⁺.turbconv.environment
    gm = state⁺
    gm_a = aux⁺

    zLL = altitude(m, aux_int)
    a_up_surf,
    θ_liq_up_surf,
    q_tot_up_surf,
    θ_liq_cv,
    q_tot_cv,
    θ_liq_q_tot_cv,
    tke = subdomain_surface_values(turbconv.surface, turbconv, m, gm, gm_a, zLL)

    @unroll_map(N_up) do i
        up[i].ρaw = FT(0)
        up[i].ρa = a_up_surf[i] * gm.ρ
        up[i].ρaθ_liq = up[i].ρa * θ_liq_up_surf[i]
        up[i].ρaq_tot = up[i].ρa * q_tot_up_surf[i]
    end
    a_en = environment_area(gm, gm_a, N_up)
    en.ρatke = gm.ρ * a_en * tke
    en.ρaθ_liq_cv = gm.ρ * a_en * θ_liq_cv
    en.ρaq_tot_cv = gm.ρ * a_en * q_tot_cv
    en.ρaθ_liq_q_tot_cv = gm.ρ * a_en * θ_liq_q_tot_cv
end;
function turbconv_boundary_state!(
    nf,
    bc::EDMFTopBC,
    m::AtmosModel{FT},
    state⁺::Vars,
    aux⁺::Vars,
    n⁻,
    state⁻::Vars,
    aux⁻::Vars,
    t,
    state_int::Vars,
    aux_int::Vars,
) where {FT}
    nothing
end;

# The boundary conditions for second-order unknowns
function turbconv_normal_boundary_flux_second_order!(
    nf,
    bc::EDMFBottomBC,
    m::AtmosModel{FT},
    fluxᵀn::Vars,
    n⁻,
    state⁻::Vars,
    diff⁻::Vars,
    hyperdiff⁻::Vars,
    aux⁻::Vars,
    state⁺::Vars,
    diff⁺::Vars,
    hyperdiff⁺::Vars,
    aux⁺::Vars,
    t,
    _...,
) where {FT}
    nothing
end;
function turbconv_normal_boundary_flux_second_order!(
    nf,
    bc::EDMFTopBC,
    m::AtmosModel{FT},
    fluxᵀn::Vars,
    n⁻,
    state⁻::Vars,
    diff⁻::Vars,
    hyperdiff⁻::Vars,
    aux⁻::Vars,
    state⁺::Vars,
    diff⁺::Vars,
    hyperdiff⁺::Vars,
    aux⁺::Vars,
    t,
    _...,
) where {FT}
    turbconv = m.turbconv
    N_up = n_updrafts(turbconv)
    up_flx = fluxᵀn.turbconv.updraft
    en_flx = fluxᵀn.turbconv.environment
    # @unroll_map(N_up) do i
    #     up_flx[i].ρaw = -n⁻ * FT(0)
    #     up_flx[i].ρa = -n⁻ * FT(0)
    #     up_flx[i].ρaθ_liq = -n⁻ * FT(0)
    #     up_flx[i].ρaq_tot = -n⁻ * FT(0)
    # end
    # en_flx.∇tke = -n⁻ * FT(0)
    # en_flx.∇e_int_cv = -n⁻ * FT(0)
    # en_flx.∇q_tot_cv = -n⁻ * FT(0)
    # en_flx.∇e_int_q_tot_cv = -n⁻ * FT(0)

end;


text:
import Flatten: flattenable
@metadata initial_value nothing
import FieldMetadata: @units, units
import FieldMetadata: @limits, limits
import FieldMetadata: @prior, prior
import FieldMetadata: @description, description
@metadata bounds nothing
import FieldMetadata: @logscaled, logscaled
@metadata reference nothing

"""
    AbstractParameters{T} <: AbstractVector{T}

An abstract type for AIBECS model parameters.

Parameters in AIBECS use the following convenience packages:

- Parameters
- FieldMetadata
- FieldDefaults
- Flatten
- Unitful
- DataFrames
- Distributions

These aim to allow for some nice features, which include

- nice syntax for unpacking parameters in functions via the `@unpack` macro (fron UnPack.jl)
- additional metadata on parameters
- easy conversion to and from vectors
- use of units and automatic conversions if necessary
- pretty table-format displays
- loading and saving to and from CSV files
- prior estimates for bayesian inference and optimization

See the list of examples to get an idea of how to generate parameters for your model.

# Examples

Generate a simple parameter type via

```jldoctest
julia> struct SimpleParams{T} <: AbstractParameters{T}
           α::T
           β::T
           γ::T
       end
SimpleParams
```

To create an instance of the `SimpleParams(Float64)` type, you can do

```jldoctest
julia> p = SimpleParams(1.0, 2.0, 3.0)
SimpleParams{Float64}
│ Row │ Symbol │ Value   │
│     │ Symbol │ Float64 │
├─────┼────────┼─────────┤
│ 1   │ α      │ 1.0     │
│ 2   │ β      │ 2.0     │
│ 3   │ γ      │ 3.0     │
```

One of the core features from Parameters is unpacking in functions, e.g.,

```jldoctest
julia> function simplef(p)
           @unpack α, γ = p
           return α + γ
       end
simplef (generic function with 1 method)

julia> simplef(p) # 1.0 + 3.0
4.0
```

More complex examples are permitted by adding metadata (thanks to FieldMetadata.jl).
You can add units

```jldoctest
julia> @units struct UnitParams{T} <: AbstractParameters{T}
           α::T | u"km"
           β::T | u"hr"
           γ::T | u"m/s"
       end ;

julia> p = UnitParams(1.0, 2.0, 3.0)
UnitParams{Float64}
│ Row │ Symbol │ Value   │ Unit     │
│     │ Symbol │ Float64 │ Unitful… │
├─────┼────────┼─────────┼──────────┤
│ 1   │ α      │ 1.0     │ km       │
│ 2   │ β      │ 2.0     │ hr       │
│ 3   │ γ      │ 3.0     │ m s^-1   │
```

Note that when adding units to your parameters, they will be converted to SI
when unpacked, as in, e.g.,

```jldoctest
julia> function speed(p)
           @unpack α, β, γ = p
           return α / β + γ
       end
speed (generic function with 1 method)

julia> speed(p) # (1.0 km / 2.0 hr + 3 m/s) in m/s
3.138888888888889
```

Another example for optimizable/flattenable parameters

```jldoctest
julia> @initial_value @units @flattenable struct OptParams{T} <: AbstractParameters{T}
           α::T | 3.6 | u"km"  | true
           β::T | 1.0 | u"hr"  | false
           γ::T | 1.0 | u"m/s" | true
       end ;

julia> p = OptParams(initial_value(OptParams)...)
OptParams{Float64}
│ Row │ Symbol │ Value   │ Initial value │ Unit     │ Optimizable │
│     │ Symbol │ Float64 │ Float64       │ Unitful… │ Bool        │
├─────┼────────┼─────────┼───────────────┼──────────┼─────────────┤
│ 1   │ α      │ 3.6     │ 3.6           │ km       │ 1           │
│ 2   │ β      │ 1.0     │ 1.0           │ hr       │ 0           │
│ 3   │ γ      │ 1.0     │ 1.0           │ m s^-1   │ 1           │
```

Thanks to the FieldMetaData interface, you can chain the following preloaded metadata:

- initial_value
- units (from Unitful.jl)
- prior (from Distributions.jl)
- description (`String`)
- bounds (2-element `Tuple`)
- logscaled (`Bool`)
- flattenable (to convert to vectors of optimizable parameters only)
- reference (`String`)

Here is an example of parameter with all the possible metadata available in AIBECS:

```jldoctest
julia> @initial_value @units @prior @description @bounds @logscaled @flattenable @reference struct FullParams{T} <: AbstractParameters{T}
           α::T | 1.0 | u"km"  | Normal(0,1)    | "The distance"   | (-Inf, Inf) | false | false | "Jean et al., 2042"
           β::T | 2.0 | u"hr"  | LogNormal(0,1) | "The time"       | (   0, Inf) | true  | true  | "Claude et al. 1983"
           γ::T | 3.0 | u"mol" | Normal(1,2)    | "The # of moles" | (  -1,   1) | false | true  | "Dusse et al. 2000"
       end ;

julia> FullParams(4.0, 5.0, 6.0)
FullParams{Float64}
│ Row │ Symbol │ Value   │ Initial value │ Unit     │ Prior                            │ Description    │ Bounds      │ Logscaled │ Optimizable │ Reference          │
│     │ Symbol │ Float64 │ Float64       │ Unitful… │ Distribu…                        │ String         │ Tuple…      │ Bool      │ Bool        │ String             │
├─────┼────────┼─────────┼───────────────┼──────────┼──────────────────────────────────┼────────────────┼─────────────┼───────────┼─────────────┼────────────────────┤
│ 1   │ α      │ 4.0     │ 1.0           │ km       │ Normal{Float64}(μ=0.0, σ=1.0)    │ The distance   │ (-Inf, Inf) │ 0         │ 0           │ Jean et al., 2042  │
│ 2   │ β      │ 5.0     │ 2.0           │ hr       │ LogNormal{Float64}(μ=0.0, σ=1.0) │ The time       │ (0, Inf)    │ 1         │ 1           │ Claude et al. 1983 │
│ 3   │ γ      │ 6.0     │ 3.0           │ mol      │ Normal{Float64}(μ=1.0, σ=2.0)    │ The # of moles │ (-1, 1)     │ 0         │ 1           │ Dusse et al. 2000  │
```

Note that there is no check that the metadata you give is consistent.
These metadata will hopefully be useful for advanced usage of AIBECS, e.g., using prior information and/or bounds for optimization.
"""
abstract type AbstractParameters{T} end
# TODO: Should it still be a subtype of AVec?
const APar = AbstractParameters # alias to make code more readable

"""
    symbols(p)

Returns the symbols in `p`.

Can also be used directly on the type of `p`
(because the symbols of `p::T` are contained in the type `T`).
"""
symbols(::Type{T}) where {T <: APar} = fieldnames(T)
symbols(::T) where {T <: APar} = symbols(T)

"""
    flattenable_symbols(p)

Returns the flattenable symbols in `p`.

The flattenable symbols are those symbols that are kepth when using `p` as a vector,
e.g., when doing `vec(p)`.
(Useful when passing parameters to an optimization routine expeting a vector of optimizable parameters.)

Can also be used directly on the type of `p`
(because the flattenable symbols of `p::T` are contained in the type `T`).
"""
flattenable_symbols(::T) where {T <: APar} = flattenable_symbols(T)
flattenable_symbols(::Type{T}) where {T <: APar} = fieldnames(T)[collect(flattenable(T))]

"""
    values(p::T) where {T <: AbstractParameters}

Returns a vector of **all** the values of `p`.

Note that `values(p)` is different from `vec(p)`.
"""
Base.values(p::T) where {T <: APar} = [getfield(p, s) for s in symbols(T)]

"""
    flattenable_values(p::T) where {T <: AbstractParameters}

Returns a vector of the **flattenable** values of `p`.

Note that `vec(p)` is different from `values(p)`.
"""
flattenable_values(p::T) where {T <: APar} = [getfield(p, s) for s in flattenable_symbols(T)]

"""
    vec(p::T) where {T <: AbstractParameters}

Returns a **SI-unit-converted** vector of flattenable values of `p`.

Note that `vec(p) ≠ flattenable_values(p)` if `p` has units.
"""
Base.vec(p::T) where {T <: APar} = [UnPack.unpack(p, Val(s)) for s in flattenable_symbols(T)]

"""
    length(p::AbstractParameter)

Returns the length of the **flattened/optimizable** vector of `p`.

May be different from the number of parameters.
Can also be used directly on the type of `p`.
"""
Base.length(::T) where {T <: APar} = sum(flattenable(T))
Base.length(::Type{T}) where {T <: APar} = sum(flattenable(T))

"""
    size(p::AbstractParameter)

Returns the size of the **flattened/optimizable** vector of `p`.

May be different from the number of parameters.
Can also be used directly on the type of `p`.
"""
Base.size(::T) where {T <: APar} = (length(T),)
Base.size(::Type{T}) where {T <: APar} = (length(T),)

function Base.show(io::IO, p::T) where T <: APar
    print(T)
    t = table(p)
    show(io, t; summary=false)
end
function Base.show(io::IO, m::MIME"text/plain", p::T) where T <: APar
    print(T)
    t = table(p)
    show(io, m, t; summary=false)
end

function table(p::T, nf::NamedTuple) where {T <: APar}
    t = DataFrame(Symbol = collect(symbols(p)), Value = values(p))
    for (n,f) in zip(keys(nf), nf)
        setproperty!(t, n, collect(f(p)))
    end
    return t
end

"""
    table(p)

Returns a `DataFrame` (a table) of `p`.
Useful for printing and saving into an actual text/latex table.
"""
function table(p::APar)
    ks, vs = Symbol[], Function[]
    all(isnothing, initial_value(p)) || (push!(ks, Symbol("Initial value")); push!(vs, initial_value))
    all(isequal(1), units(p)) || (push!(ks, :Unit); push!(vs, units))
    all(isnothing, prior(p)) || (push!(ks, Symbol("Prior")); push!(vs, prior))
    all(isnothing, limits(p)) || (push!(ks, Symbol("Limits")); push!(vs, limits))
    all(isempty, description(p)) || (push!(ks, Symbol("Description")); push!(vs, description))
    all(isnothing, bounds(p)) || (push!(ks, Symbol("Bounds")); push!(vs, bounds))
    any(logscaled(p)) && (push!(ks, Symbol("Logscaled")); push!(vs, logscaled))
    all(flattenable(p)) || (push!(ks, Symbol("Optimizable")); push!(vs, flattenable))
    all(isnothing, reference(p)) || (push!(ks, Symbol("Reference")); push!(vs, reference))
    nf = (; zip(ks, vs)...)
    t = table(p, nf)
end
# These line prevent the default (0.0, 1.0) value
limits(::T) where {T<:APar} = limits(T)
limits(::Type{T}) where {T<:APar} = Tuple(nothing for _ in 1:fieldcount(T))

"""
    latex(p)

Returns a LaTeX-formatted table of the parameters.
"""
latex(p::APar) = show(stdout, MIME("text/latex"), table(p))

"""
    unpack(p <: AbstractParameters, s)

Unpacks the parameter `s` from `p`.

Note this is specialized and will convert the parameter value to SI units.
"""
@inline UnPack.unpack(p::T, ::Val{f}) where {T<:APar,f} = ustrip(upreferred(getproperty(p, f) * units(T, f)))

"""
    +(p::T, v::Vector) where {T <: AbstractParameters}

Adds the flattened vector `v` to `p`.

**Warning:** This method for `+` is implemented only for differentiation
using dual and hyperdual numbers.
If you want to change the values of `p`, you should do so explicitly
rather than use this `+` method.
"""
Base.:+(p::T, v::Vector) where {T <: APar} = reconstruct(T.name.wrapper, flattenable_values(p) + v)

"""
    reconstruct(T, v)

Reconstructs the parameter of type `T` from flattenable vector `v`.

Note that `T` must be the generic type (or type wrapper if you will).
I.e., `T` must *not* include the parametric scalar type, as in,
use `T == Params` instead of `T == Params{Float64}`.
"""
function reconstruct(::Type{T}, v::Vector{V}) where {T <: APar, V}
    all(isnothing, initial_value(T)) && error("Can't reconstruct without initial values")
    reconstructed_v = convert(Vector{V}, collect(initial_value(T))) # fill in initial values
    reconstructed_v[collect(flattenable(T))] .= v            # fill in v
    return T{V}(reconstructed_v...)  # wrapper allows to change type
end

"""
    getindex(p::T, i) where {T <: AbstractParameters}

Returns the i-th element of vec(p).

This is not efficient and only used for testing the derivatives with ForwardDiff.
"""
Base.getindex(p::T, i) where {T <: APar} = flattenable_values(p)[i]

function (::Type{T})(args::Quantity...) where {T <: APar}
    all(isequal(1), units(T)) && error("$T needs `units` for this construction to work")
    return T([ustrip(x |> units(T, f)) for (x,f) in zip(args, fieldnames(T))]...)
end

function (::Type{T})(;kwargs...) where {T <: APar}
    all(isnothing, initial_value(T)) && length(kwargs) ≠ length(symbols(T)) && error("$T needs `initial_value` if one of the parameters is not supplied")
    value(f::Symbol, v::Quantity) = ustrip(v |> units(T, f))
    value(f::Symbol, v) = v
    return T([f ∈ keys(kwargs) ? value(f, kwargs[f]) : initial_value(T, f) for f in fieldnames(T)]...)
end

#===============================
Writing to savable formats
===============================#

function Base.Dict(p::T, s=symbols(p)) where {T<:APar}
    v = [getfield(p,s) for s in s]
    u = [units(p,s) for s in s]
    return Dict([(s,v*u) for (s,v,u) in zip(s,v,u)])
end
function Base.NamedTuple(p::T, s=symbols(p)) where {T<:APar}
    v = [getfield(p,s) for s in s]
    u = [units(p,s) for s in s]
    return (; zip(s, v .* u)...)
end

#=====================
mismatch of parameters
=====================#

# in p space (should not be needed except for checking/testing)
mismatch(p::T, k::Symbol) where {T<:APar} = mismatch(T, bijector(T, k)(p))
mismatch(p::T) where {T<:APar} = mismatch(T, p2λ(p))
mismatch(::Type{T}, p::Tp) where {T<:APar, Tp<:APar} = mismatch(p)

# in λ space
mismatch(::Type{T}, λ) where {T<:APar} = sum(mismatch(T, k, λᵢ) for (k, λᵢ) in zip(flattenable_symbols(T), λ))
mismatch(::Type{T}, k, λ) where {T<:APar} = mismatch(prior(T, k), λ)
mismatch(d::ContinuousUnivariateDistribution, λ) = -logpdf(transformed(d), λ)

# TODO check if I need that below:
# The functions below is just for ForwardDiff to work with vectors instead of p
# which requires `mismatch` to know about the priors, which are containted in `T`
function generate_objective(ωs, μx, σ²x, v, ωp, ::Type{T}) where {T<:APar}
    nt, nb = length(ωs), length(v)
    tracers(x) = state_to_tracers(x, nb, nt)
    f(x, p) = ωp * mismatch(T, p) +
        sum([ωⱼ * mismatch(xⱼ, μⱼ, σⱼ², v) for (ωⱼ, xⱼ, μⱼ, σⱼ²) in zip(ωs, tracers(x), μx, σ²x)])
    return f
end

#==================
Change of variables
==================#

# using Bijectors and evaluating parameter mismatch in λ-space
# Usage should look like this
# F = AIBECSFunction(...)
# f, ∇ₓf = generate_objective(...) # This requires to dispatch them like f(p::APar) = ... and f(λ::Vector) = ...
# obj, grad, hess = F1Method(F, ∇ₓF, f, ∇ₓf, ...) # directly in λ-space
# p2λ = bijector(MyParams) # just reexport bijector and inverse
# λ2p = inverse(p2λ)
# λ₀ = p2λ(p₀)
# results = optimize(obj, grad, hess, λ₀, ...)
# λ_opt = results.minimizer
# p_opt = λ2p(λ_opt) # and done and dusted...
# some actual code that I should write and test
import Bijectors: bijector
"""
    bijector(T::AbstractParameters [, k::Symbol])

Returns the function for the change of variables of the parameters.

The function is a bijection from the supports/domains of the priors to ℝⁿ,
from the Bijectors.jl package.
You can specify the parameter symbol to get the bijector of that parameter.
"""
bijector(::Type{T}) where {T<:APar} = bijector.([p for p in prior(T) if !isnothing(p)])
bijector(::Type{T}, k) where {T<:APar} = bijector(prior(T, k))
export bijector

"""
    p2λ(p::AbstractParameters)

Converts `p` to a real-valued vector for optimization.
(referred to as the λ space in AIBECS)
"""
p2λ(p::T) where {T<:APar} = [bijector(T,k)(v) for (k,v) in zip(flattenable_symbols(p), flattenable_values(p))]
"""
    λ2p(T::Type{AbstractParameters}, λ::Vector)

Converts real-valued vector `λ` back to parameters object `p`.

Note that the instance of your parameters type `T` is required here because
it contains information on non-optimizable parameters and priors of optimizable parameters
"""
function λ2p(::Type{T}, λs) where {T<:APar}
    ps = [inv(bijector(T,k))(λ) for (k,λ) in zip(flattenable_symbols(T), λs)]
    reconstruct(T, ps)
end
export p2λ, λ2p

export AbstractParameters, latex


text:
"""
    get_locations(csv::String) -> Vector{Int32}

Extract the values of the `Location` column from the CSV file.

# Arguments
- `csv::String`: $(CSV_ARGUMENT)

# Returns
- `Vector{Int32}`: the locations

# Throws
- Exceptions from: [`get_all`](@ref)

# Example
```jldoctest; output = false
using PDFHighlights
HEADER = PDFHighlights.Internal.Constants.HEADER

_file, io = mktemp()
println(io, HEADER, '\\n', ",,,,5722")
flush(io)
file = _file * ".csv"
mv(_file, file)

get_locations(file) == Int32[5722]

# output

true
```
"""
function get_locations(csv::String)::Vector{Int32}
    _, _, _, _, locations = get_all(csv)
    return locations
end


text:
# This file is auto-generated by AWSMetadata.jl
using AWS
using AWS.AWSServices: sso
using AWS.Compat
using AWS.UUIDs

"""
    get_role_credentials(account_id, role_name, x-amz-sso_bearer_token)
    get_role_credentials(account_id, role_name, x-amz-sso_bearer_token, params::Dict{String,<:Any})

Returns the STS short-term credentials for a given role name that is assigned to the user.

# Arguments
- `account_id`: The identifier for the AWS account that is assigned to the user.
- `role_name`: The friendly name of the role that is assigned to the user.
- `x-amz-sso_bearer_token`: The token issued by the CreateToken API call. For more
  information, see CreateToken in the AWS SSO OIDC API Reference Guide.

"""
get_role_credentials(account_id, role_name, x_amz_sso_bearer_token; aws_config::AbstractAWSConfig=global_aws_config()) = sso("GET", "/federation/credentials", Dict{String, Any}("account_id"=>account_id, "role_name"=>role_name, "headers"=>Dict{String, Any}("x-amz-sso_bearer_token"=>x_amz_sso_bearer_token)); aws_config=aws_config)
get_role_credentials(account_id, role_name, x_amz_sso_bearer_token, params::AbstractDict{String, <:Any}; aws_config::AbstractAWSConfig=global_aws_config()) = sso("GET", "/federation/credentials", Dict{String, Any}(mergewith(_merge, Dict{String, Any}("account_id"=>account_id, "role_name"=>role_name, "headers"=>Dict{String, Any}("x-amz-sso_bearer_token"=>x_amz_sso_bearer_token)), params)); aws_config=aws_config)

"""
    list_account_roles(account_id, x-amz-sso_bearer_token)
    list_account_roles(account_id, x-amz-sso_bearer_token, params::Dict{String,<:Any})

Lists all roles that are assigned to the user for a given AWS account.

# Arguments
- `account_id`: The identifier for the AWS account that is assigned to the user.
- `x-amz-sso_bearer_token`: The token issued by the CreateToken API call. For more
  information, see CreateToken in the AWS SSO OIDC API Reference Guide.

# Optional Parameters
Optional parameters can be passed as a `params::Dict{String,<:Any}`. Valid keys are:
- `"max_result"`: The number of items that clients can request per page.
- `"next_token"`: The page token from the previous response output when you request
  subsequent pages.
"""
list_account_roles(account_id, x_amz_sso_bearer_token; aws_config::AbstractAWSConfig=global_aws_config()) = sso("GET", "/assignment/roles", Dict{String, Any}("account_id"=>account_id, "headers"=>Dict{String, Any}("x-amz-sso_bearer_token"=>x_amz_sso_bearer_token)); aws_config=aws_config)
list_account_roles(account_id, x_amz_sso_bearer_token, params::AbstractDict{String, <:Any}; aws_config::AbstractAWSConfig=global_aws_config()) = sso("GET", "/assignment/roles", Dict{String, Any}(mergewith(_merge, Dict{String, Any}("account_id"=>account_id, "headers"=>Dict{String, Any}("x-amz-sso_bearer_token"=>x_amz_sso_bearer_token)), params)); aws_config=aws_config)

"""
    list_accounts(x-amz-sso_bearer_token)
    list_accounts(x-amz-sso_bearer_token, params::Dict{String,<:Any})

Lists all AWS accounts assigned to the user. These AWS accounts are assigned by the
administrator of the account. For more information, see Assign User Access in the AWS SSO
User Guide. This operation returns a paginated response.

# Arguments
- `x-amz-sso_bearer_token`: The token issued by the CreateToken API call. For more
  information, see CreateToken in the AWS SSO OIDC API Reference Guide.

# Optional Parameters
Optional parameters can be passed as a `params::Dict{String,<:Any}`. Valid keys are:
- `"max_result"`: This is the number of items clients can request per page.
- `"next_token"`: (Optional) When requesting subsequent pages, this is the page token from
  the previous response output.
"""
list_accounts(x_amz_sso_bearer_token; aws_config::AbstractAWSConfig=global_aws_config()) = sso("GET", "/assignment/accounts", Dict{String, Any}("headers"=>Dict{String, Any}("x-amz-sso_bearer_token"=>x_amz_sso_bearer_token)); aws_config=aws_config)
list_accounts(x_amz_sso_bearer_token, params::AbstractDict{String, <:Any}; aws_config::AbstractAWSConfig=global_aws_config()) = sso("GET", "/assignment/accounts", Dict{String, Any}(mergewith(_merge, Dict{String, Any}("headers"=>Dict{String, Any}("x-amz-sso_bearer_token"=>x_amz_sso_bearer_token)), params)); aws_config=aws_config)

"""
    logout(x-amz-sso_bearer_token)
    logout(x-amz-sso_bearer_token, params::Dict{String,<:Any})

Removes the client- and server-side session that is associated with the user.

# Arguments
- `x-amz-sso_bearer_token`: The token issued by the CreateToken API call. For more
  information, see CreateToken in the AWS SSO OIDC API Reference Guide.

"""
logout(x_amz_sso_bearer_token; aws_config::AbstractAWSConfig=global_aws_config()) = sso("POST", "/logout", Dict{String, Any}("headers"=>Dict{String, Any}("x-amz-sso_bearer_token"=>x_amz_sso_bearer_token)); aws_config=aws_config)
logout(x_amz_sso_bearer_token, params::AbstractDict{String, <:Any}; aws_config::AbstractAWSConfig=global_aws_config()) = sso("POST", "/logout", Dict{String, Any}(mergewith(_merge, Dict{String, Any}("headers"=>Dict{String, Any}("x-amz-sso_bearer_token"=>x_amz_sso_bearer_token)), params)); aws_config=aws_config)


text:
function reinterpret(a::Array{MultiValue{S,T,N,L}}) where {S,T,N,L}
  b = reinterpret(T,a)
  sa = size(a)
  sv = Size(S)
  t = _Size_to_tuple(sv)
  s = (t...,sa...)
  reshape(b,s)
end

_Size_to_tuple(::Size{t}) where t = t


text:
function s2riskmargin(p::ProjParam, scr, coc)
  invs = InvPort(p.t_0, p.T, p.cap_mkt, p.invs_par...)
  proj = Projection(p.tax_rate, p.tax_credit_0, p.cap_mkt, invs,
                    p.l_ins, p.l_other, p.dyn)
  bal_vec  = vcat(proj.val_0, proj.val)
  discount = 1 ./ cumprod(1 .+ p.cap_mkt.rfr.x)
  tp = Vector( bal_vec[!,:tpg] +
                      bal_vec[!,:bonus] +
                      bal_vec[!,:cost_prov])[1:(p.T-p.t_0)]
  src_future = (tp * scr / tp[1])
  coc * src_future ⋅ discount
end

"""
`s2bal(p::ProjParam)`

Calculation of the unshocked S2 balance sheet
"""
function s2bal(p::ProjParam)
  invs = InvPort(p.t_0, p.T, p.cap_mkt, p.invs_par...)
  proj = Projection(p.tax_rate, p.tax_credit_0, p.cap_mkt, invs,
                    p.l_ins, p.l_other, p.dyn)
  bal =
    hcat(proj.val_0, DataFrame(tax_credit = proj.tax_credit_0))
  return hcat(bal, DataFrame(bof = bof(bal), scen = :be))
end

"""
`s2bal(p::ProjParam, md::S2Module, shock!::Any, scen::Symbol)`

Calculation of the shocked S2 balance sheet for module `md`
"""
function s2bal(p::ProjParam,
               md::S2Module,
               shock!::Any,
               scen::Symbol)
  cpm = deepcopy(p.cap_mkt)
  l_ins = deepcopy(p.l_ins)
  if md.shock_object ==  :CapMkt shock!(cpm, md) end
  invs = InvPort(p.t_0, p.T, cpm, p.invs_par...)
  if md.shock_object == :CapMkt_AdjVal0 shock!(invs, md) end
  if md.shock_object == :LiabIns shock!(l_ins, md) end
  if md.shock_object == :InvPort_LiabIns
    shock!(invs, l_ins, md)
  end
  proj = Projection(p.tax_rate, p.tax_credit_0, cpm, invs,
                    l_ins, p.l_other, p.dyn)
  if md.shock_object == :CapMkt_AdjVal0
    mkt_val0_adj!(proj, invs, md, scen)
  end
  bal =
    hcat(proj.val_0, DataFrame(tax_credit = proj.tax_credit_0))
  return hcat(bal, DataFrame(bof = bof(bal), scen = scen))
end

"""
`bof(bal::DataFrame)`

Helper function for the calculation of basic own funds
"""
bof(bal::DataFrame) =
  bal[1, :invest][1,1] +
  bal[1, :tax_credit][1,1] -
  bal[1, :tpg][1,1] -
  bal[1, :cost_prov][1,1] -
  bal[1, :bonus][1,1]

"""
`bof(md::S2Module, scen::Symbol)`

Basic own funds for module `md` and scenario `scen`
"""
bof(md::S2Module, scen::Symbol) =
  bof(md.balance[md.balance[!,:scen] .== scen, :])

"""
`fdb(md::S2Module, scen::Symbol)`

Future discretionary benefits for module `md` and scenario `scen`
"""
fdb(md::S2Module, scen::Symbol) =
  md.balance[md.balance[!,:scen] .== scen, :bonus][1,1]

"""
`scr!(mdl::S2Module)`

Scenario based SCR calculation for `mdl::S2Module`

**Changed:** `mdl::S2Module`
"""
function scr!(mdl::S2Module)
  shock_keys = collect(keys(mdl.shock))
  net =
    (bof(mdl, :be) .- Float64[bof(mdl, 𝑠𝑚) for 𝑠𝑚 ∈ shock_keys])
  gross =
    (net .+ fdb(mdl, :be) .-
     Float64[fdb(mdl, 𝑠𝑚) for 𝑠𝑚 ∈ shock_keys])
  if :corr in fieldnames(typeof(mdl))
    mdl.scr[NET] = sqrt(net ⋅ (mdl.corr * net))
    mdl.scr[GROSS] = sqrt(gross ⋅ (mdl.corr * gross))
  else
    i = findmax(net)[2]
    mdl.scr[NET] = max(0, net[i])
    mdl.scr[GROSS] = max(0, gross[i])
  end
end

"""
`scr(md::S2Module, corr::Matrix{Float64})`

Aggregation of SCRs of sub-modules
"""
function scr(md::S2Module, corr::Matrix{Float64})
  _scr = zeros(Float64, 2)
  net = Float64[md.mds[𝑖].scr[NET] for 𝑖 ∈ 1:length(md.mds)]
  gross = Float64[md.mds[𝑖].scr[GROSS] for 𝑖 ∈ 1:length(md.mds)]
  _scr[GROSS] = sqrt(gross ⋅ (corr * gross))
  _scr[NET] = sqrt(net ⋅ (corr * net))
  return _scr
end

## S2MktInt -----------------------------------------------------
"""
`scr!(mkt_int::S2MktInt)`

SCR for interest rate risk, sub-module `mkt_int`

**Changed:** `mkt_int::S2MktInt`
"""
function scr!(mkt_int::S2MktInt)
  shock_keys = collect(keys(mkt_int.shock))
  net =
    bof(mkt_int, :be) .-
  Float64[bof(mkt_int, 𝑠𝑚) for 𝑠𝑚 ∈ shock_keys]
  gross =
    net .+ fdb(mkt_int, :be) -
    Float64[fdb(mkt_int, 𝑠𝑚) for 𝑠𝑚 ∈ shock_keys]

  i_up = findfirst( (in)([:spot_up]), shock_keys)
  i_down = findfirst( (in)([:spot_down]), shock_keys)

  mkt_int.scen_up = net[i_up] >= net[i_down]
  mkt_int.scr[NET] = maximum([0.0; net])
  mkt_int.scr[GROSS] =
    max(0.0, mkt_int.scen_up ? gross[i_up] : gross[i_down])
end

"""
`rfrshock(rfr::Vector{Float64}, s2_mkt_int, int_type)`

Helper function: shock for the risk free interest rate
"""
function rfrshock(rfr::Vector{Float64}, s2_mkt_int, int_type)
  ## shock the risk free interest rate
  len = min(length(rfr),
            length(s2_mkt_int.shock[:spot_up]),
            length(s2_mkt_int.shock[:spot_down]))
  spot = forw2spot(rfr[1:len])
  if int_type == :spot_down
    forw =
      spot2forw(spot .*
                (1 .+ s2_mkt_int.shock[:spot_down][1:len]))
  elseif int_type == :spot_up
    forw =
      spot2forw(spot .+
                max.(spot .* s2_mkt_int.shock[:spot_up][1:len],
                    s2_mkt_int.spot_up_abs_min))
  else # :be
      forw = spot2forw(spot)
  end
  return forw
end

"""
`mktintshock!(cap_mkt::CapMkt, s2_mkt_int, int_type::Symbol)`

Shock for interest rate market risk

**Changed:** `cap_mkt::CapMkt`
"""
function mktintshock!(cap_mkt::CapMkt,
                      s2_mkt_int,
                      int_type::Symbol)
  cap_mkt.rfr.x =
    deepcopy(rfrshock(cap_mkt.rfr.x, s2_mkt_int, int_type))
end

## S2MktEq ------------------------------------------------------
"""
`mkteqshock!(invs::InvPort, mkt_eq, eq_type::Symbol)`

Shock for equity market risk

**Changed:** `invs::InvPort`
"""
function mkteqshock!(invs::InvPort, mkt_eq, eq_type::Symbol)
  for 𝑖𝑛𝑣𝑒𝑠𝑡 ∈ invs.igs[:IGStock].investments
    if mkt_eq.eq2type[𝑖𝑛𝑣𝑒𝑠𝑡.name] == eq_type
      𝑖𝑛𝑣𝑒𝑠𝑡.proc.x *= (1 + mkt_eq.shock[eq_type])
    end
  end
end

"""
`mkt_val0_adj!(proj::Projection, invs::InvPort,
               mkt_eq, eq_type::Symbol)`

Adjust initial market value for S2 balance sheet in order to
reflect that the initial fall of the market value of equity
investments is reflected in the shockedsolvency balance sheet
Must be called after the projection

**Changed:** `proj::Projection`
"""
function mkt_val0_adj!(proj::Projection, invs::InvPort,
                       mkt_eq, eq_type::Symbol)
  for 𝑖𝑛𝑣𝑒𝑠𝑡 ∈ invs.igs[:IGStock].investments
    if mkt_eq.eq2type[𝑖𝑛𝑣𝑒𝑠𝑡.name] == eq_type
      proj.val_0[1,:invest] +=
        mkt_eq.shock[eq_type] * 𝑖𝑛𝑣𝑒𝑠𝑡.mv_0
    end
  end
end

## S2Def1 -------------------------------------------------------
"""
`scr!(def::S2Def1)`

SCR for interest rate risk, sub-module `def::S2Def1`

**Changed:** `def::S2Def1`
"""
function scr!(def::S2Def1)
  var = def.tlgd ⋅ (def.u * def.tlgd) + def.v ⋅ def.slgd
  sigma_norm = -sqrt(var)/sum(def.tlgd)
  if sigma_norm <= def.scr_par[:low][1]
    def.scr[NET] = def.scr_par[:low][2] * sqrt(var)
  elseif sigma_norm <= def.scr_par[:medium][1]
    def.scr[NET] = def.scr_par[:medium][2] * sqrt(var)
  else
    def.scr[NET] = sum(def.tlgd)
  end
  def.scr[GROSS] = def.scr[NET]
end

## S2LifeBio ----------------------------------------------------
"""
`selectmort!(p::ProjParam, bio::S2LifeBio)`

Identify those model points that are subject to mortality
risk. This function does not properly take into account
second order effects due to the effect of boni.
However, for realistic portfolios second order effects are
unlikely to change the set of identified model points.

**Changed:** `bio::S2LifeBio  (bio.mp_select)`
"""
function selectmort!(p::ProjParam, bio::S2LifeBio)
  #invs = InvPort(p.t_0, p.T, p.cap_mkt, p.invs_par...)
  for 𝑠𝑦𝑚𝑏 ∈ collect(keys(bio.shock))
    merge!(bio.mp_select,
           Dict(𝑠𝑦𝑚𝑏 => Array{Bool}(undef, length(p.l_ins.mps))))
    for (𝑚, 𝑚𝑝) ∈ enumerate(p.l_ins.mps)
      if (𝑠𝑦𝑚𝑏 == :sx_mass_pension) & (!𝑚𝑝.pension_contract)
        bio.mp_select[𝑠𝑦𝑚𝑏][𝑚] = false
      else
        tp = tpg(p.t_0,
                 p.cap_mkt.rfr.x,
                 𝑚𝑝)
        mp_shock = deepcopy(𝑚𝑝)
        bioshock!(mp_shock, bio, 𝑠𝑦𝑚𝑏)
        tp_shock = tpg(p.t_0,
                       p.cap_mkt.rfr.x,
                       mp_shock)
        bio.mp_select[𝑠𝑦𝑚𝑏][𝑚] = (tp_shock > tp)
      end
    end
  end
end

"""
`bioshock!(mp::ModelPoint, bio::S2LifeBio, symb::Symbol)`

Helper function: shock for biometric risk

**Changed:** `mp::ModelPoint`
"""
function bioshock!(mp::ModelPoint, bio::S2LifeBio, symb::Symbol)
  if symb in [:qx, :px]
    qxpxshock!(mp, bio, symb)
  elseif symb in [:sx_down, :sx_up,
                  :sx_mass_pension, :sx_mass_other]
    sxshock!(mp, bio, symb)
  elseif symb in [:cat]
    catshock!(mp, bio, symb)
  end
end

"""
`bioshock!(l_ins::LiabIns, bio::S2LifeBio, shock_symb::Symbol)`

Shock for biometric risk

**Changed:** `l_ins::LiabIns  (l_ins.mps)`
"""
function bioshock!(l_ins::LiabIns,
                   bio::S2LifeBio,
                   shock_symb::Symbol)
  for (𝑚, 𝑚𝑝) ∈ enumerate(l_ins.mps)
    if bio.mp_select[shock_symb][𝑚]
      bioshock!(𝑚𝑝, bio, shock_symb)
    end
  end
end

"""
`qxpxshock!(mp::ModelPoint, bio::S2LifeBio, symb::Symbol)`

Helper function: shock for mortality risk

**Changed:** `mp::ModelPoint`
"""
function qxpxshock!(mp::ModelPoint, bio::S2LifeBio, symb::Symbol)
  mp.prob[!,:qx] =
    min.(1, (1 + bio.shock[symb]) * Vector( mp.prob[!,:qx]))
  mp.prob[!,:sx] = min.(1 .- mp.prob[!,:qx], mp.prob[!,:sx])
  mp.prob[!,:px] =  1.0 .- mp.prob[!,:qx] - mp.prob[!,:sx]
end

"""
`sxshock!(mp::ModelPoint, bio::S2LifeBio, symb::Symbol)`

Helper function: shock for surrender risk

**Changed:** `mp::ModelPoint`
"""
function sxshock!(mp::ModelPoint, bio::S2LifeBio, symb::Symbol)
  if symb == :sx_down
    mp.prob[!,:sx] =
      max.((1 + bio.shock[:sx_down]) * Vector( mp.prob[!,:sx]),
          Vector(
                  mp.prob[!,:sx]) .+
                  bio.shock_param[:sx_down_threshold])
  elseif symb == :sx_up
    mp.prob[!,:sx] =
      min.(1, (1 + bio.shock[symb]) * Vector( mp.prob[!,:sx]))
  elseif symb == :sx_mass_pension
    mp.prob[1, :sx] = bio.shock[symb]
  elseif symb == :sx_mass_other
    mp.prob[1, :sx] = bio.shock[symb]
  end
  mp.prob[!,:qx] = min.(1 .- mp.prob[!,:sx], mp.prob[!,:qx])
  mp.prob[!,:px] =  1.0 .- mp.prob[!,:qx] - mp.prob[!,:sx]
end

"""
`catshock!(mp::ModelPoint, bio::S2LifeBio, symb::Symbol)`

Shock for mortality catastrophe risk

**Changed:** `mp::ModelPoint`
"""
function catshock!(mp::ModelPoint, bio::S2LifeBio, symb::Symbol)
  mp.prob[1, :qx] = min(1, mp.prob[1, :qx] + bio.shock[symb])
  mp.prob[1, :sx] = min(1 .- mp.prob[1, :qx], mp.prob[1, :sx])
  mp.prob[!,:px] =  1.0 .- mp.prob[!,:qx] - mp.prob[!,:sx]
end

## S2LifeCost ---------------------------------------------------
"""
`costshock!(invs::InvPort, l_ins::LiabIns, cost::S2LifeCost)`

Shock for expense risk

**Changed:** `invs::InvPort`, `l_ins::LiabIns`
"""
function costshock!(invs::InvPort,
                    l_ins::LiabIns,
                    cost::S2LifeCost)
  shock_eoy =
    (1 + cost.shock[:cost]) *
    (1 + cost.shock_param[:infl]) .^ collect(1:l_ins.dur)
  for 𝑠𝑦𝑚𝑏 ∈ collect(keys(invs.igs))
    invs.igs[𝑠𝑦𝑚𝑏].cost.rel .*= shock_eoy
    invs.igs[𝑠𝑦𝑚𝑏].cost.abs .*= shock_eoy
  end
  for 𝑚𝑝 ∈ l_ins.mps
    𝑚𝑝.λ[:, :boy] *= (1 + cost.shock[:cost])
    𝑚𝑝.λ[:, :eoy] *= (1 + cost.shock[:cost])
    𝑚𝑝.λ[:, :infl] .+= cost.shock_param[:infl]
    𝑚𝑝.λ[:, :cum_infl] =
      𝑚𝑝.λ[1, :cum_infl] / (1 + 𝑚𝑝.λ[1, :infl]) *
      cumprod(1 .+ 𝑚𝑝.λ[:, :infl])
  end
end

## S2Op ---------------------------------------------------------
"""
`scr!(op::S2Op, bscr)`

SCR for interest rate risk, sub-module `op`

**Changed:** `op`
"""
function scr!(op::S2Op, bscr)
  ## SCR for operational risk
  op.comp_prem =
    op.fac[:prem] *
    (op.prem_earned +
       max(0,
           op.prem_earned -
             op.fac[:prem_py] * op.prem_earned_prev))
  op.comp_tp = op.fac[:tp]  * max(0, op.tp)
  op.scr =
    min(op.fac[:bscr] * bscr, max(op.comp_prem, op.comp_tp)) +
    op.fac[:cost] * op.cost_ul
end

## S2 -----------------------------------------------------------
"""
`scr!(s2::S2, tax_credit_0::Float64)`

Total SCR

**Changed:** `s2`
"""
function scr!(s2::S2, tax_credit_0::Float64)
  s2.bscr = scr(s2, s2.corr)
  scr!(s2.op, s2.bscr[GROSS])
  s2.adj_tp =
    -max(0.0, min(s2.bscr[GROSS] - s2.bscr[NET], fdb(s2, :be)))
  s2.adj_dt =
    -max(tax_credit_0 - (s2.bscr[GROSS] + s2.op.scr + s2.adj_tp),
         0)
  s2.scr = s2.bscr[GROSS] + s2.adj_tp + s2.adj_dt + s2.op.scr
end


text:
export Ensemble, ensemble_sample_psdd, sample_vtree, learn_ensemble_llw!, learn_ensemble_em!,
       learn_ensemble_stacking!

using ..Utils: sample_vtree, kfold
using ThreadPools

"Weighted ensemble of probabilistic circuits."
mutable struct Ensemble{T <: ProbCircuit}
    C::Vector{T}
    W::Vector{Float64}
end

"""
Creates an ensemble of `n` SamplePSDD-generated probabilistic circuits, with `v` the total number
of variables in the data, `ϕ` the logic constraints, `D` the data to be learned and `k` the
maximum number of primes to be sampled.

Keyword arguments for `sample_psdd` are passed down. Optionally, the function takes keyword
argument `vtree_bias`, which samples more left (value closer to 0.0) or right-leaning (value closer
to 1.0) vtrees. If a negative value is given, sample uniformly distributed vtrees.

Weights are computed by the given `strategy`. These can be any one of the following:
    1. `:likelihood` for likelihood weighting;
    2. `:uniform` for uniform weights;
    3. `:em` for Expectation-Maximization;
    4. `:stacking` for mixture model Stacking;
"""
function ensemble_sample_psdd(n::Integer, ϕ::Bdd, k::Int, D::DataFrame; vtree_bias::Real = -1.0,
        strategy::Symbol = :em, verbose::Bool = true, em_maxiter::Integer = 100,
        kfold::Integer = min(nrow(D), 5), pseudocount::Real = 1.0, kwargs...)::Ensemble{StructProbCircuit}
    circs = Vector{StructProbCircuit}(undef, n)
    v = ncol(D)
    verbose && print("Sampling ", n, " PSDDs...\n  ")
    @qthreads for i ∈ 1:n
        V = sample_vtree(v, vtree_bias)
        circs[i] = sample_psdd(ϕ, V, k, D; pseudocount, kwargs...)
        verbose && print('*')
    end
    verbose && print('\n')
    c = log(n)
    E = Ensemble{StructProbCircuit}(circs, -fill(c, n))
    if strategy == :likelihood return learn_ensemble_llw!(E, D)
    elseif strategy == :em return learn_ensemble_em!(E, D; verbose, maxiter = em_maxiter)
    elseif strategy == :stacking
        return learn_ensemble_stacking!(E, D; verbose, maxiter = em_maxiter, k = kfold, pseudocount)
    end
    @assert strategy == :uniform "Unrecognized ensemble strategy."
    return E
end

"Learns the weights of the Ensemble by the likelihood value of data `D`."
function learn_ensemble_llw!(E::Ensemble{T}, D::DataFrame)::Ensemble{T} where T <: ProbCircuit
    n = length(E.C)
    LL = Vector{Float64}(undef, n)
    @qthreads for i ∈ 1:n
        @inbounds LL[i] = sum(EVI(E, D))
    end
    W = exp.(LL .- maximum(LL))
    E.W .= log.(W ./ sum(W))
    return E
end

"Learns the weights of the Ensemble by Expectation-Maximization."
function learn_ensemble_em!(E::Ensemble{T}, D::DataFrame; maxiter::Integer = 100,
        reuse::Union{Matrix{Float64}, Nothing} = nothing,
        verbose::Bool = true)::Ensemble{T} where T <: ProbCircuit
    N, K = nrow(D), length(E.C)
    ln_N = log(N)
    W = Matrix{Float64}(undef, N, K)
    N_k = Vector{Float64}(undef, K)
    ll = Vector{Float64}(undef, N)
    if isnothing(reuse)
        verbose && println("Pre-computing component log-likelihoods...")
        LL = Matrix{Float64}(undef, N, K)
        @qthreads for i ∈ 1:K @inbounds LL[:,i] .= log_likelihood_per_instance(E.C[i], D) end
    else LL = reuse end
    for j ∈ 1:maxiter
        @qthreads for i ∈ 1:K @inbounds W[:,i] .= E.W[i] .+ LL[:,i] end
        Threads.@threads for i ∈ 1:N @inbounds W[i,:] .-= logsumexp(W[i,:]) end
        @qthreads for i ∈ 1:K @inbounds N_k[i] = logsumexp(W[:,i]) end
        E.W .= N_k .- ln_N
        @qthreads for i ∈ 1:K @inbounds W[:,i] .= LL[:,i] .+ E.W[i] end
        Threads.@threads for i ∈ 1:N @inbounds ll[i] = logsumexp(W[i,:]) end
        verbose && println("EM Iteration ", j, "/", maxiter, ". Log likelihood ", sum(ll))
    end
    return E
end

"Learns the weights of the Ensemble by Stacking, with `k` as the number of folds in k-fold."
function learn_ensemble_stacking!(E::Ensemble{T}, D::DataFrame; maxiter::Integer = 100,
        k::Integer = min(nrow(D), 5), pseudocount::Real = 1.0,
        verbose::Bool = true)::Ensemble{T} where T <: ProbCircuit
    N, K = nrow(D), length(E.C)
    F = kfold(N, k)
    LL = Matrix{Float64}(undef, N, K)
    for j ∈ 1:k
        I, J = F[j]
        D_T, D_R = D[I,:], D[J,:]
        @qthreads for i ∈ 1:K estimate_parameters(E.C[i], D_R; pseudocount = pseudocount == 0 ? 1.0 : pseudocount) end
        @qthreads for i ∈ 1:K LL[I,i] .= log_likelihood_per_instance(E.C[i], D_T) end
        verbose && println("Stacking fold ", j, '/', k, '.')
    end
    learn_ensemble_em!(E, D; maxiter, verbose, reuse = LL)
    @qthreads for i ∈ 1:K estimate_parameters(E.C[i], D; pseudocount) end
    return E
end

function weighted_query(E::Ensemble{T}, D::DataFrame, f::Function; kwargs...)::Vector{Float64} where T <: ProbCircuit
    n, k = nrow(D), length(E.C)
    M = Matrix{Float64}(undef, n, k)
    @qthreads for j ∈ 1:k
        ll = f(E.C[j], D; kwargs...)
        @inbounds ll .+= E.W[j]
        for i ∈ 1:n @inbounds M[i,j] = ll[i] end
    end
    return logsumexp.(eachrow(M))
end

@inline function log_likelihood_per_instance(E::Ensemble{T}, D::DataFrame;
        use_gpu::Bool = false)::Vector{Float64} where T <: ProbCircuit
    return weighted_query(E, D, log_likelihood_per_instance; use_gpu)
end

@inline function marginal(E::Ensemble{T}, D::DataFrame)::Vector{Float64} where T <: ProbCircuit
    return weighted_query(E, D, marginal)
end


text:
m = Model()

@variable(m, x)
@variable(m, y)
@variable(m, z)

xu = UnitJuMP.UnitVariableRef(x, u"km")
yu = UnitJuMP.UnitVariableRef(y, u"hr")
zu = UnitJuMP.UnitVariableRef(z, u"s")

@test -xu == UnitJuMP.UnitAffExpr(-x, u"km")
@test xu + 200u"m" == UnitJuMP.UnitAffExpr(x + 0.2, u"km")
@test 200u"m" + xu  == UnitJuMP.UnitAffExpr(1000x + 200, u"m")
@test xu - 400u"m" == UnitJuMP.UnitAffExpr(x - 0.4, u"km")
@test 400u"m" - xu  == UnitJuMP.UnitAffExpr(-1000x + 400, u"m")

@test 200u"km" * x == UnitJuMP.UnitAffExpr(200 * x, u"km")
@test 200 * xu == UnitJuMP.UnitAffExpr(200 * x, u"km")
@test 1.5u"s" * xu == UnitJuMP.UnitAffExpr(1.5 * x, u"s*km")

@test xu / 0.5 == UnitJuMP.UnitAffExpr(2 * x, u"km")
@test x / 0.5u"s" == UnitJuMP.UnitAffExpr(2 * x, u"s^-1")
@test xu / 0.5u"s" == UnitJuMP.UnitAffExpr(2 * x, u"km/s")

@test yu + zu == UnitJuMP.UnitAffExpr(y + z/3600, u"hr")
@test yu - zu == UnitJuMP.UnitAffExpr(y - z/3600, u"hr")

expr = yu + 1800 * zu

@test expr + 3600u"s" == UnitJuMP.UnitAffExpr(y + 0.5z + 1, u"hr")
@test expr - 1800u"s" == UnitJuMP.UnitAffExpr(y + 0.5z - 0.5, u"hr")
@test 2u"s" + expr == UnitJuMP.UnitAffExpr(3600y + 1800z + 2, u"s")
@test 3600u"s" - expr == UnitJuMP.UnitAffExpr(-3600y -1800z + 3600, u"s")

@test 2 * expr == UnitJuMP.UnitAffExpr(2y + z, u"hr")
@test expr / 0.5 == UnitJuMP.UnitAffExpr(2y + z, u"hr")
@test 2u"kW" * expr == UnitJuMP.UnitAffExpr(2y + z, u"kW*hr")
@test expr / 0.5u"km" == UnitJuMP.UnitAffExpr(2y + z, u"hr/km")

@variable(m, w)
wu = UnitJuMP.UnitVariableRef(w, u"minute")

@test expr + wu == UnitJuMP.UnitAffExpr(y + 0.5z + w/60, u"hr")
@test expr - wu == UnitJuMP.UnitAffExpr(y + 0.5z - w/60, u"hr")
@test wu + expr == UnitJuMP.UnitAffExpr(60y + 30z + w, u"minute")
@test wu - expr == UnitJuMP.UnitAffExpr(-60y - 30z + w, u"minute")
@test expr + yu == UnitJuMP.UnitAffExpr(2y + 0.5z, u"hr")

expr2 = 30wu + yu

@test expr + expr2 == UnitJuMP.UnitAffExpr(2y + 0.5z + 0.5w, u"hr")
@test expr2 + expr == UnitJuMP.UnitAffExpr(120y + 30z + 30w, u"minute")
@test 2*expr - expr2 == UnitJuMP.UnitAffExpr(y + z -0.5w, u"hr")


text:
using OmegaCore
using OmegaTest
using Distributions
using Test

@testset "multivariate" begin
  μ = 5
  xs = 1 ~ Mv(Normal(5, 1), (1000,))
  samples = randsample(xs)
  @test mean(samples) ≈ μ atol = 0.3
  @test isinferred(randsample, xs)
end

@testset "Multivariate 2" begin
  x = 1 ~ Normal(0, 1)
  function f(id, ω)
    x(ω) + Uniform(0, 1)(id, ω)
  end
  xs = 2 ~ Mv(f, (3, 3))
  randsample((x, xs))
end