### A Pluto.jl notebook ###
# v0.14.1

using Markdown
using InteractiveUtils

# ╔═╡ 63ba08cc-59a8-11eb-0a0f-27efac60d779
using Pkg, DrWatson

# ╔═╡ 6db218c6-59a8-11eb-2a8b-7107354cf590
begin
	#@quickactivate "StatisticalRethinkingStan"
	using StanSample, StanOptimize
	using StatisticalRethinking
end

# ╔═╡ 51fc19b8-59a8-11eb-2214-15aca59b807b
md" ## Clip-08-01-06s.jl"

# ╔═╡ 8aaa4bcc-59a8-11eb-2003-f1213b116565
begin
	df = CSV.read(sr_datadir("rugged.csv"), DataFrame)
	dropmissing!(df, :rgdppc_2000)
	dropmissing!(df, :rugged)
	df.log_gdp = log.(df[:, :rgdppc_2000])
	df.log_gdp_s = df.log_gdp / mean(df.log_gdp)
	df.rugged_s = df.rugged / maximum(df.rugged)
	data = (N = size(df, 1), G = df.log_gdp_s,
		R = df.rugged_s)
	PRECIS(df[:, [:log_gdp, :log_gdp_s, :rugged, :rugged_s]])
end

# ╔═╡ d7d3e626-6a45-11eb-1820-a3ec98e556b1
stan8_1_1 = "
parameters {
	real a;
	real b;
}

model {
	a ~ normal(1, 1);
	b ~ normal(0, 1);
}
";

# ╔═╡ a3e7c070-6a46-11eb-072b-3943db854020
begin
	m8_1_1s = SampleModel("m8.1.1s", stan8_1_1)
	rc8_1_1s = stan_sample(m8_1_1s)
	if success(rc8_1_1s)
		post8_1_1s_df = read_samples(m8_1_1s, :dataframe)
		PRECIS(post8_1_1s_df[:, [:a, :b]])
	end
end

# ╔═╡ ac6a235c-6a67-11eb-0d6c-c38eb2162598
begin
	x = 0:0.01:1
	p1 = plot(;ylim=(0.6, 1.4), leg=false)
	for i in 1:20:4000
		y = post8_1_1s_df.a[i] .+ post8_1_1s_df.b[i] .* x
		plot!(x, y, color=:lightgrey)
	end
	plot!(x, mean(post8_1_1s_df.a) .+ mean(post8_1_1s_df.b) .* x, color=:darkblue)
end;

# ╔═╡ 7f74d204-6a6a-11eb-3931-4de1dfb8d527
stan8_1_2 = "
parameters {
	real a;
	real b;
}

model {
	a ~ normal(1, 0.1);
	b ~ normal(0, 0.1);
}
";

# ╔═╡ 8f79ecfc-6a6a-11eb-1173-2fa6e4609d72
begin
	m8_1_2s = SampleModel("m8.1.2s", stan8_1_2)
	rc8_1_2s = stan_sample(m8_1_2s)
	if success(rc8_1_2s)
		post8_1_2s_df = read_samples(m8_1_2s, :dataframe)
		PRECIS(post8_1_2s_df[:, [:a, :b]])
	end
end

# ╔═╡ baaabbf4-6a6a-11eb-091d-cd8a0e80004f
begin
	p2 = plot(;ylim=(0.6, 1.4), leg=false)
	for i in 1:80:4000
		y = post8_1_2s_df.a[i] .+ post8_1_2s_df.b[i] .* x
		plot!(x, y, color=:lightgrey)
	end
	plot!(x, mean(post8_1_2s_df.a) .+ mean(post8_1_2s_df.b) .* x, color=:darkblue)

	plot(p1, p2, layout=(1, 2))
end

# ╔═╡ 9a39b1c0-6a6c-11eb-05ce-ab82a8d58dfe
stan8_1_3 = "
parameters {
	real a;
	real b;
real sigma;
}

model {
	a ~ normal(1, 0.1);
	b ~ normal(0, 0.3);
sigma ~ exponential(1);
}
";

# ╔═╡ 6269a87e-6a6c-11eb-2e14-39df73ddd6d8
begin
	m8_1_3s = SampleModel("m8.1.3s", stan8_1_3)
	rc8_1_3s = stan_sample(m8_1_3s; data)
	if success(rc8_1_3s)
		post8_1_3s_df = read_samples(m8_1_3s, :dataframe)
		PRECIS(post8_1_3s_df[:, [:a, :b, :sigma]])
	end
end

# ╔═╡ 6b462290-6a6c-11eb-2d9b-97241415b825
begin
	plot(;ylim=(0.6, 1.4), leg=false)
	for i in 1:80:4000
		y = post8_1_3s_df.a[i] .+ post8_1_3s_df.b[i] .* x
		plot!(x, y, color=:lightgrey)
	end
	plot!(x, mean(post8_1_3s_df.a) .+ mean(post8_1_3s_df.b) .* x, color=:darkblue)
end

# ╔═╡ e890c59c-6a68-11eb-0d33-4d52a21d0ccf
stan8_1_4 = "
data {
	int N;
	vector[N] G;
	vector[N] R;
}

parameters {
	real a;
	real b;
	real<lower=0> sigma;
}

transformed parameters {
	vector[N] mu;
	mu = a + b * (R - 0.125);
}

model {
	a ~ normal(1, 1);
	b ~ normal(0, 1);
	sigma ~ exponential(1);
	G ~ normal(mu, sigma);
}
";

# ╔═╡ 8e815d5c-6a66-11eb-21dc-99734c31f4e1
begin
	m8_1_4s = SampleModel("m8.1.4s", stan8_1_4)
	rc8_1_4s = stan_sample(m8_1_4s; data)
	if success(rc8_1_4s)
		post8_1_4s_df = read_samples(m8_1_4s, :dataframe)
		PRECIS(post8_1_4s_df[:, [:a, :b, :sigma]])
	end
end

# ╔═╡ 04a39568-6a6c-11eb-19ea-f55be3914a51
begin
	plot(;ylim=(0.6, 1.4), leg=false)
	for i in 1:80:4000
		y = post8_1_4s_df.a[i] .+ post8_1_4s_df.b[i] .* x
		plot!(x, y, color=:lightgrey)
	end
	plot!(x, mean(post8_1_4s_df.a) .+ mean(post8_1_4s_df.b) .* x, color=:darkblue)
end

# ╔═╡ 45767e2e-6a63-11eb-3e12-354a7e32a374
md" ## End of clip-08-01-06s.jl"

# ╔═╡ Cell order:
# ╟─51fc19b8-59a8-11eb-2214-15aca59b807b
# ╠═63ba08cc-59a8-11eb-0a0f-27efac60d779
# ╠═6db218c6-59a8-11eb-2a8b-7107354cf590
# ╠═8aaa4bcc-59a8-11eb-2003-f1213b116565
# ╠═d7d3e626-6a45-11eb-1820-a3ec98e556b1
# ╠═a3e7c070-6a46-11eb-072b-3943db854020
# ╠═ac6a235c-6a67-11eb-0d6c-c38eb2162598
# ╠═7f74d204-6a6a-11eb-3931-4de1dfb8d527
# ╠═8f79ecfc-6a6a-11eb-1173-2fa6e4609d72
# ╠═baaabbf4-6a6a-11eb-091d-cd8a0e80004f
# ╠═9a39b1c0-6a6c-11eb-05ce-ab82a8d58dfe
# ╠═6269a87e-6a6c-11eb-2e14-39df73ddd6d8
# ╠═6b462290-6a6c-11eb-2d9b-97241415b825
# ╠═e890c59c-6a68-11eb-0d33-4d52a21d0ccf
# ╠═8e815d5c-6a66-11eb-21dc-99734c31f4e1
# ╠═04a39568-6a6c-11eb-19ea-f55be3914a51
# ╟─45767e2e-6a63-11eb-3e12-354a7e32a374


@enum CameraKind Perspective Orthographic Pixel

# import GeometryBasics: Vec, Mat
@component @with_kw mutable struct Camera3D
    lookat::Vec3f0             = zero(Vec3f0)
    up::Vec3f0                 = Z_AXIS
    right::Vec3f0              = X_AXIS
    fov::Float32               = 42.0f0
    near::Float32              = 0.1f0
    far::Float32               = 3000.0f0
    view                       :: Mat4f0
    proj                       :: Mat4f0
    projview                   :: Mat4f0
    rotation_speed::Float32    = 0.001f0
    translation_speed::Float32 = 0.02f0
    locked::Bool               = false
    camerakind::CameraKind     = Perspective
end

function Camera3D(width_pixels::Integer, height_pixels::Integer; eyepos = -10 * Y_AXIS,
                  lookat = zero(Vec3f0), up = Z_AXIS, right = X_AXIS, near = 0.1f0,
                  far = 3000.0f0, fov = 42.0f0, camerakind = Perspective)
    up    = normalizeperp(lookat - eyepos, up)
    right = normalize(cross(lookat - eyepos, up))

    viewm = lookatmat(eyepos, lookat, up)
    zoom  = norm(forward(eyepos, lookat))
    projm = projmat(width_pixels, height_pixels, zoom, near, far, fov, camerakind)
    return Camera3D(; lookat = lookat, up = up, right = right, fov = fov, near = near,
                    far = far, view = viewm, proj = projm, projview = projm * viewm,
                    camerakind = camerakind)
end
function projmat(w::Integer, h::Integer, cam::Camera3D, zoom)
    return projmat(w, h, zoom, cam.near, cam.far, cam.fov, cam.camerakind)
end

function projmat(w::Integer, h::Integer, zoom, near, far, fov, camerakind)
    aspect = w / h
    h = tan(fov / 360.0f0 * pi) * near
    w = h * aspect
    if camerakind == Pixel
        return diagm(ones(Float32, 4))
    elseif camerakind == Orthographic
        h, w = h * zoom * 23.0f0, w * zoom * 23.0f0
        return projmatortho(Float32, -w, w, -h, h, near, far)
    elseif camerakind == Perspective
        return projmatpersp(Float32, fov, aspect, near, far)
    end
end

abstract type InteractiveSystem <: System end

struct CameraOperator <: InteractiveSystem end

Overseer.requested_components(::CameraOperator) = (Spatial, Camera3D, Canvas)

function Overseer.update(::CameraOperator, m::AbstractLedger)
    canvas = singleton(m, Canvas)
    mouse = singleton(m, Mouse)
    keyboard = singleton(m, Keyboard)

    x, y         = mouse.x, mouse.y
    mouse_button = canvas.mouse_buttons
    w, h         = size(canvas)

    scroll_dx, scroll_dy = mouse.dscroll
    dx = mouse.dx
    dy = mouse.dy

    @inbounds for e in @entities_in(m, Spatial && Camera3D)
        e.locked && continue
        cam = e[Camera3D]
        new_pos = Point3f0(e.position)
        #world orientation/mouse stuff
        new_lookat = cam.lookat
        if pressed(mouse) && !(keyboard.button ∈ CTRL_KEYS && pressed(keyboard))
            if mouse.button == GLFW.MOUSE_BUTTON_1 #rotation
                trans1  = translmat(-cam.lookat)
                rot1    = rotate(dy * cam.rotation_speed, -cam.right)
                rot2    = rotate(-dx * cam.rotation_speed, cam.up)
                trans2  = translmat(cam.lookat)
                mat_    = trans2 * rot2 * rot1 * trans1
                new_pos = Point3f0((mat_*Vec4f0(new_pos..., 1.0f0))[1:3])

            elseif mouse.button == GLFW.MOUSE_BUTTON_2 #panning
                rt         = cam.right * dx * cam.translation_speed / 2
                ut         = -cam.up * dy * cam.translation_speed / 2
                new_lookat -= rt + ut
                new_pos    -= rt + ut
            end
        end

        #keyboard stuff
        if pressed(keyboard)
            if keyboard.button in WASD_KEYS
                new_pos, cam.lookat = wasd_event(new_pos, cam, keyboard)
            end
        end

        #scroll stuff no dxp
        new_forward = forward(new_pos, cam.lookat)
        new_pos += Point3f0(new_forward * scroll_dy * cam.translation_speed * 5)
        zoom = norm(new_forward)
        cam.proj = projmat(w, h, cam, zoom)

        # update_viewmat
        u_forward = normalize(new_forward)
        cam.right = unitright(u_forward, cam.up)
        cam.up = unitup(u_forward, e.right)
        cam.view = lookatmat(new_pos, cam.lookat, cam.up)
        cam.projview = cam.proj * cam.view
        e[Spatial] = Spatial(new_pos, e.velocity)
    end
end

unitforward(position, lookat) = normalize(forward(position, lookat))
unitright(forward, up)        = normalize(right(forward, up))
unitup(forward, right)        = normalize(up(forward, right))

forward(position, lookat) = Vec3f0(lookat - position)
right(forward, up)        = cross(forward, up)
up(forward, right)        = cross(right, forward)

#maybe it would be better to make the eyepos up etc vectors already vec4 but ok
function wasd_event(position::Point3f0, cam::Camera3D, keyboard)
    #ok this is a bit hacky but fine
    # origlen = norm(position)
    button = keyboard.button
    new_lookat = cam.lookat
    if button == GLFW.KEY_A
        move = Point3f0(cam.translation_speed * 5 * cam.right)
        #the world needs to move in the opposite direction
        position   -= move
        new_lookat -= move
    end
    if button == GLFW.KEY_D
        move       = Point3f0(cam.translation_speed * 5 * cam.right)
        position   += move
        new_lookat += move
    end
    if button == GLFW.KEY_W
        move       = Point3f0(unitforward(position, cam.lookat) * 5 * cam.translation_speed)
        position   += move
        new_lookat += move
    end
    if button == GLFW.KEY_S
        move       = Point3f0(unitforward(position, cam.lookat) * 5 * cam.translation_speed)
        position   -= move
        new_lookat -= move
    end
    return position, new_lookat
end


using GLVisualize, GLAbstraction

if !isdefined(:runtests)
    window = glscreen()
    timesignal = loop(linspace(0f0,1f0,360))
end

function juliadata(max_iterations, imgx, imgy)
    scalex, scaley = 4.0/imgx, 4.0/imgy

    # initialize our Float32 heightfield
    heightfield = zeros(Float32, imgx, imgy)

    # do julia set magic!
    for x=1:imgx, y=1:imgy
        cy = y * scaley - 2.0
        cx = x * scalex - 2.0
        z = Complex(cx, cy)
        c = Complex(-0.4, 0.6)
        i = 0
        for t in 0:max_iterations
            if norm(z) > 2.0
                break
            end
            z = z * z + c
            i = t
        end
        heightfield[x,y] = -(i/512f0)
    end
    heightfield
end

rotation_angle = const_lift(*, timesignal, 2f0*pi)
rotation = map(rotationmatrix_z, rotation_angle)

heightfield = juliadata(256, 700, 700)

# visualize the heightfield as a surface
vis = visualize(
    heightfield, :surface,
    model=rotation
)

# display it on the window
_view(vis, window)

if !isdefined(:runtests)
    renderloop(window) # render!
end


using TropicalTensors
using ForwardDiff
using Test

@testset "spinglass" begin
    lt = Cylinder(10, 8)
    sg = Spinglass(lt, ones(Float32, 152), zeros(Float32, 80))
    res = solve(sg; usecuda=false)
    @test res.n == 152
end

@testset "counting tropical" begin
    lt = Cylinder(10, 8)
    sg = Spinglass(lt, ones(Float32, 152), zeros(Float32, 80))
    res = solve(CountingTropicalF32, sg; usecuda=false)
    @test res.n == 152
    @test res.c == 2
end

using BitBasis
function count_degeneracy_exact(sg::Spinglass{LT,T}) where {LT, T}
    elist = T[]
    bonds = sgbonds(sg.lattice)
    for b in basis(BitStr64{length(sg.lattice)})
        eng = zero(T)
        for ((i, j), J) in zip(bonds, sg.Js)
            eng += b[i] == b[j] ? J : -J
        end
        push!(elist, eng)
    end
    findall(==(minimum(elist)), elist) |> length
end

@testset "counting tropical 2" begin
    lt = Cylinder(4, 4)
    sg = Spinglass(lt, rand([-1, 1], 28), zeros(Int, 16))
    res = solve(CountingTropical{Int,Int}, sg; usecuda=false)
    @test res.c == count_degeneracy_exact(sg)
end

@testset "forwarddiff" begin
    L = 3
    Js = ones(Float32, 2*L^2-L)
    gs = ForwardDiff.gradient(x->solve(Cylinder(L, L), x, zeros(eltype(x), L^2); usecuda=false).n, Js)
    @test gs ≈ ones(Float32, 2*L^2-L)
end


const PWGRID_VERSION = 3

mutable struct GVectorsW
    Ngwx::Int64
    idx_gw2r::Array{Int}
end

mutable struct GVectors
    Ng::Int64
    G::Array{Float64,2}
    G2::Array{Float64}
end

mutable struct PWGrid
    ecutwfc::Float64
    ecutrho::Float64
    Ns::Array{Int64}
    LatVecs::Array{Float64,2}
    RecVecs::Array{Float64,2}
    Ω::Float64
    r::Array{Float64,2}
    gvec::GVectors
    gvecw::GVectorsW
end

function PWGrid( ecutwfc::Float64, LatVecs::Array{Float64,2} )
    ecutrho = 4.0*ecutwfc
    #
    RecVecs = 2*pi*inv(LatVecs')
    Ω = det(LatVecs)
    #
    LatVecsLen = Array{Float64}(undef,3)
    LatVecsLen[1] = norm(LatVecs[:,1])
    LatVecsLen[2] = norm(LatVecs[:,2])
    LatVecsLen[3] = norm(LatVecs[:,3])

    Ns = Array{Int64}(undef,3)
    Ns[1] = 2*round( Int, sqrt(ecutrho/2)*LatVecsLen[1]/pi ) + 1
    Ns[2] = 2*round( Int, sqrt(ecutrho/2)*LatVecsLen[2]/pi ) + 1
    Ns[3] = 2*round( Int, sqrt(ecutrho/2)*LatVecsLen[3]/pi ) + 1

    # Use even sampling numbers
    Ns[1] = Ns[1] % 2 == 1 ? Ns[1] + 1 : Ns[1]
    Ns[2] = Ns[2] % 2 == 1 ? Ns[2] + 1 : Ns[2]
    Ns[3] = Ns[3] % 2 == 1 ? Ns[3] + 1 : Ns[3]

    Npoints = prod(Ns)
    r = init_grid_R( Ns, LatVecs )

    gvec = init_grid_G( Ns, RecVecs )
    gvecw = init_gvecw( ecutwfc, gvec.G2 )

    return PWGrid( ecutwfc, ecutrho, Ns, LatVecs, RecVecs, Ω, r, gvec, gvecw )
end

function mm_to_nn(mm::Int,S::Int)
    if mm > S/2
        return mm - S
    else
        return mm
    end
end

function init_grid_G( Ns, RecVecs )

    Ng = prod(Ns)

    G  = Array{Float64}(undef,3,Ng)
    G2 = Array{Float64}(undef,Ng)

    ig = 0
    for k in 0:Ns[3]-1
    for j in 0:Ns[2]-1
    for i in 0:Ns[1]-1
        ig = ig + 1
        gi = mm_to_nn( i, Ns[1] )
        gj = mm_to_nn( j, Ns[2] )
        gk = mm_to_nn( k, Ns[3] )
        G[1,ig] = RecVecs[1,1]*gi + RecVecs[1,2]*gj + RecVecs[1,3]*gk
        G[2,ig] = RecVecs[2,1]*gi + RecVecs[2,2]*gj + RecVecs[2,3]*gk
        G[3,ig] = RecVecs[3,1]*gi + RecVecs[3,2]*gj + RecVecs[3,3]*gk
        G2[ig] = G[1,ig]^2 + G[2,ig]^2 + G[3,ig]^2
    end
    end
    end

    return GVectors( Ng, G, G2 )
end

function init_gvecw( ecutwfc, G2 )
    idx_gw2r = findall( 0.5*G2 .< ecutwfc )
    Ngwx = length(idx_gw2r)
    return GVectorsW( Ngwx, idx_gw2r )
end

function find_edges( Ns )
    eS = Ns/2 .+ 0.5
    edges = []
    ip = 1
    for k in 0:Ns[3]-1
    for j in 0:Ns[2]-1
    for i in 0:Ns[1]-1
        ei = abs(i - eS[1])
        ej = abs(j - eS[2])
        ek = abs(k - eS[3])
        # if any of i, j, or k is equal to Ns/2 or Ns/2+1
        if ei < 1.0 || ej < 1.0 || ek < 1.0
            push!(edges,ip)
        end
        ip = ip + 1
    end
    end
    end
    return edges
end

function init_grid_R( Ns, LatVecs )
    #
    Npoints = prod(Ns)
    #
    R = Array{Float64}(undef,3,Npoints)
    ip = 0
    for k in 0:Ns[3]-1
    for j in 0:Ns[2]-1
    for i in 0:Ns[1]-1
        ip = ip + 1
        R[1,ip] = LatVecs[1,1]*i/Ns[1] + LatVecs[1,2]*j/Ns[2] + LatVecs[1,3]*k/Ns[3]
        R[2,ip] = LatVecs[2,1]*i/Ns[1] + LatVecs[2,2]*j/Ns[2] + LatVecs[2,3]*k/Ns[3]
        R[3,ip] = LatVecs[3,1]*i/Ns[1] + LatVecs[3,2]*j/Ns[2] + LatVecs[3,3]*k/Ns[3]
    end
    end
    end
    #
    return R
end


# test_direct_discretization.jl

# direct discretization for elliptic problem in 2D/3D
function get_problem_2d(n)
    A = MultigridMethod((n,m) -> p₂₂(n), (n, n), V(3, 2))
    b = fill(1,(n-1)*(n-1))
    return (A,b)
end

@testset "Direct discretization of elliptic PDE, 2D" begin
    for n in [16 32 64 128 256 512 1024]
        (A, b) = get_problem_2d(n)
        A.grids[1].b .= b # copy rhs
        for item in Base.Iterators.take(A,15) end # iterate
        @test A.resnorm[end] < 1/n^2
        log(A,2)
    end
end

function get_problem_3d(n)
    A = MultigridMethod((n,m,k) -> p₂₃(n), (n, n, n), V(3, 2))
    b = fill(1,(n-1)*(n-1)*(n-1))
    return (A,b)
end

@testset "Direct discretization of elliptic PDE, 2D" begin
    for n in [4 8 16 32 64 128]
        (A, b) = get_problem_3d(n)
        A.grids[1].b .= b # copy rhs
        for item in Base.Iterators.take(A,15) end # iterate
        @test A.resnorm[end] < 1/n^3
        log(A,3)
    end
end


using DynamicBounds

x0(p) = [34.0; 20.0; 0.0; 0.0; 16.0; 0.0]
function f!(du, u, p, t)
    du[1] = -p[1]*u[1]*u[2] + p[2]*u[3] + p[6]*u[6]
    du[2] = -p[1]*u[1]*u[2] + p[2]*u[3] + p[3]*u[3]
    du[3] =  p[1]*u[1]*u[2] - p[2]*u[3] - p[3]*u[3]
    du[4] =  p[3]*u[3] - p[4]*u[4]*u[5] + p[5]*u[6]
    du[5] = -p[4]*u[4]*u[5] + p[5]*u[6] + p[6]*u[6]
    du[6] =  p[4]*u[4]*u[5] - p[5]*u[6] - p[6]*u[6]
    return
end

tspan = (0.0,18.0e-5*250)
pL = [0.1; 0.033; 16.0; 5.0; 0.5; 0.3]
pU = 10.0*pL

prob = ODERelaxProb(f!, tspan, x0, pL, pU)

A = [0.0 -1.0 -1.0  0.0  0.0  0.0;
     0.0  0.0  0.0  0.0 -1.0 -1.0;
     1.0 -1.0 0.0 1.0 -1.0 0.0]
b = [-20.0; -16.0; -2.0]
set!(prob, PolyhedralConstraint(A, b))

xL = zeros(6)
xU = [34.0; 20.0; 20.0; 34.0; 16.0; 16.0]
set!(prob, ConstantStateBounds(xL,xU))

integrator = DifferentialInequality(prob, calculate_relax = false, calculate_subgradient = false)

relax!(integrator)

ratio = rand(6)
pstar = pL.*ratio .+ pU.*(1.0 .- ratio)
setall!(integrator, ParameterValue(), pstar)
integrate!(integrator)


using JuLIPMaterials: Vec3, Mat3, Ten33, Ten43
using JuLIPMaterials.CLE: onb3D, euclidean, spherical

using Einsum, StaticArrays

using LinearAlgebra

export GreenFunction, IsoGreenFcn3D, grad

# ========== Implementation of the 3D Green's Function =============

abstract type AbstractGreenFunction{T} end

struct GreenFunction3D{T} <: AbstractGreenFunction{T}
   Nquad::Int
   C::Ten43{T}
   remove_singularity::Bool
end

"""
`GreenFunction3D`
construct a green's function type
"""
function GreenFunction(C::Ten43; Nquad = nothing, remove_singularity = true)
   if Nquad == nothing
      error("still need to implement auto-tuning, please provide Nquad")
   end
   return GreenFunction3D(Nquad, C, remove_singularity)
end

GreenFunction(C::Array{T, 4}; kwargs...) where {T} =
      GreenFunction(Ten43{T}(C); kwargs...)

GreenFunction(at::AbstractAtoms; kwargs...) =
      GreenFunction(calculator(at), at; kwargs...)

GreenFunction(calc::AbstractCalculator, at::AbstractAtoms; kwargs...) =
      GreenFunction(elastic_moduli(calc, at); kwargs...)

function (G::GreenFunction3D{T})(x) where T
   if G.remove_singularity && norm(x) < 1e-10
      return @SMatrix zeros(T, 3, 3)
   end
   return eval_green(Vec3(x), G.C, G.Nquad)
end

function grad(G::GreenFunction3D{T}, x) where T
   if G.remove_singularity && norm(x) < 1e-10
      return @SArray zeros(T, 3, 3, 3)
   end
   return grad_green(Vec3(x), G.C, G.Nquad)
end

"eval_green(x::Vec3, ℂ::Ten43, Nquad::Int)"
function eval_green(x::Vec3{TT}, ℂ::Ten43, Nquad::Int) where TT
   # allocate
   G = @SMatrix zeros(TT, 3, 3)
   zz = @MMatrix zeros(TT, 3, 3)
   # Initialise tensors.
   x̂ = x/norm(x)
   # two vectors orthogonal to x.
   x1, x2 = onb3D(x̂)
   # Integrate
   for ω in range(0.0, step=pi/Nquad, length=Nquad)
      z = cos(ω) * x1 + sin(ω) * x2
      @einsum zz[i,j] = z[α] * ℂ[i,α,j,β] * z[β]
      # Perform integration
      G += inv(zz)
   end
   # Normalise appropriately
   return G / (4*pi*norm(x)*Nquad)
end

function grad_green(x::Vec3{TT}, ℂ::Ten43, Nquad::Int) where TT
   # allocate
   DG = @MArray zeros(TT, 3, 3, 3)
   zz = @MMatrix zeros(TT, 3, 3)
   zT = @MMatrix zeros(TT, 3, 3)
   # Initialise tensors.
   x̂ = x/norm(x)
   # two vectors orthogonal to x.
   x1, x2 = onb3D(x̂)
   # Integrate
   for ω in range(0.0, step=pi/Nquad, length=Nquad)
      z = cos(ω) * x1 + sin(ω) * x2
      @einsum zz[i,j] = z[α] * ℂ[i,α,j,β] * z[β]
      @einsum zT[i,j] = z[α] * ℂ[i,α,j,β] * x̂[β]
      zzinv = inv(zz)
      F = zzinv * (zT + zT') * zzinv
      @einsum DG[i,j,k] = DG[i,j,k] + zzinv[i,j] * x̂[k] - F[i,j] * z[k]
   end
   DG ./= (-4.0 * pi * norm(x)^2 * Nquad)
   return SArray(DG)
end

struct IsoGreenFcn3D{T} <: AbstractGreenFunction{T}
   λ::T
   μ::T
   remove_singularity::Bool
end

IsoGreenFcn3D(λ, μ; remove_singularity = true) =
   IsoGreenFcn3D(λ, μ, remove_singularity)

function (G::IsoGreenFcn3D{T})(x) where T
   if G.remove_singularity && norm(x) < 1e-10
      return @SMatrix zeros(T, 3, 3)
   end
   return eval_greeniso(Vec3(x), G.λ, G.μ)
end

function grad(G::IsoGreenFcn3D{T}, x) where T
   if G.remove_singularity && norm(x) < 1e-10
      return @SArray zeros(T, 3, 3, 3)
   end
   return grad_greeniso(Vec3{T}(x), G.λ, G.μ)
end

"isotropic CLE Green's function"
function eval_greeniso(x::Vec3{T}, λ::Real, μ::Real) where T
   Id = SMatrix{3,3,T}(I)
   x̂ = x/norm(x)
   return (((λ+3*μ)/(λ+2*μ)/norm(x)) * Id  +
               ((λ+μ)/(λ+2*μ)/norm(x)) * x̂ * x̂') / (8.0*π*μ)
end

"gradient of isotropic CLE Green's function"
function grad_greeniso(x::Vec3{T}, λ::Real, μ::Real) where T
   DG = @MArray zeros(T, 3, 3, 3)
   Id = SMatrix{3,3,T}(I)
   x̂ = x / norm(x)
   for i = 1:3, j = 1:3, k = 1:3
      DG[i,j,k] = (λ+μ) * (Id[i,k] * x̂[j] + Id[j,k] * x̂[i]) - (λ+3*μ) * Id[i,j] * x̂[k] - 3*(λ+μ) * x̂[i] * x̂[j] * x̂[k]
   end
   DG ./= 8 * π * μ*(λ+2*μ) * norm(x)^2
   return DG
end


function KS_solve_ChebySCF( pw::PWGrid,
                            V_ionic, Focc, Nstates::Int64;
                            cheby_degree=5,
                            β = 0.5, E_NN = nothing,
                            rho0 = nothing,
                            Potentials0 = nothing,
                            NiterMax=100,
                            verbose=false )

    Ngwx = pw.gvecw.Ngwx
    Ns = pw.Ns
    Npoints = prod(Ns)
    ΔV = pw.Ω/Npoints

    # starting rho
    if rho0 == nothing
        srand(1234)
        v = rand( ComplexF64, Ngwx, Nstates )
        v = ortho_gram_schmidt( v )
        rho = calc_rho( pw, Focc, v )
    else
        srand(1234)
        v = rand( ComplexF64, Ngwx, Nstates ) # What if v0 is provided ?
        rho = copy( rho )
    end

    # starting Potentials
    if Potentials0 == nothing
        # initialize potential: Ionic, Hartree, XC
        V_Hartree = real( G_to_R( Ns, Poisson_solve(pw, rho) ) )
        V_xc = excVWN( rho ) + rho .* excpVWN( rho )
        Potentials = PotentialsT( V_ionic, V_Hartree, V_xc )
    else
        Potentials = PotentialsT( Potentials0.Ionic,
                                  Potentials0.Hartree,
                                  Potentials0.XC )
    end

    Etot_old = 0.0
    rho_old = copy(rho)

    if( E_NN != nothing )
        Energies = EnergiesT( 0.0, 0.0, 0.0, 0.0, 0.0, E_NN )
    else
        Energies = EnergiesT( 0.0, 0.0, 0.0, 0.0, 0.0, 0.0 )
    end

    # eigenvalues
    λ = zeros(Nstates)
    resnrm = zeros(Nstates)

    for iter = 1:NiterMax

        ub, lb = get_ub_lb_lanczos(pw, Potentials, Nstates*2)
        #@printf("lb = %f, ub = %f\n", lb, ub)

        v = chebyfilt(pw, Potentials, v, cheby_degree, lb, ub)
        v = ortho_gram_schmidt(v)

        #
        rho_new = calc_rho( pw, Focc, v )
        diffRho = norm(rho_new - rho)
        #
        rho = β*rho_new[:] + (1-β)*rho[:]
        #
        #integRho = sum(rho)*ΔV
        #@printf("integRho = %18.10f\n", integRho)
        #
        V_Hartree = real( G_to_R( Ns, Poisson_solve(pw, rho) ) )
        V_xc = excVWN( rho ) + rho .* excpVWN( rho )
        Potentials = PotentialsT( V_ionic, V_Hartree, V_xc )

        # Calculate energies and update potetials
        Energies = calc_Energies( pw, Potentials, Focc, v, Energies.NN )

        Etot = Energies.Total
        diffE = abs( Etot - Etot_old )

        #
        @printf("ChebySCF: %8d %18.10f %18.10e %18.10e\n", iter, Etot, diffE, diffRho )

        if diffE < 1e-7
            @printf("ChebySCF is is converged: iter: %d , diffE = %10.7e\n", iter, diffE)
            break
        end
        #
        Etot_old = Etot
    end

    # Calculate eigenvalues
    Hv = op_H(pw,Potentials,v)
    G = v' * Hv
    R = Hv - v*G
    λ = real(eigvals(G))
    #for j = 1:Nstates
    #   resnrm[j] = norm(R[:,j]);
    #   @printf("λ[%2d] = %11.3e, resnrm = %11.3e\n", j, λ[j], resnrm[j]);
    #end

    return Energies, Potentials, v, λ

end


using HiQGA.transD_GP, PyPlot
## model fixed parts, i.e., air, but only to be compatible with AEM ...
zfixed   = [-1e5]
ρfixed   = [1e12]
nmax = 200
# z grid spec starts, the first z and first ρ will be unused in MT
zstart = 0.0
extendfrac, dz = 1.22, 15
zall, znall, zboundaries = transD_GP.setupz(zstart, extendfrac, dz=dz, n=50, showplot=true, atol=1e-3)
z, ρ, nfixed = transD_GP.makezρ(zboundaries; zfixed=zfixed, ρfixed=ρfixed)
## COPROD from Occam paper
T = [28.5 38.5 52.0 70.5 95.5 129.0 174.6 236.2 319.6 432.5 585.1 791.7 1071.1 1449.2 1960.7][:]
d_log10_ρ = [2.315  2.254  2.229  2.188  2.18  2.162  2.151  2.208  2.194  2.299  2.338  2.42  2.405  2.308  2.397][:]
d_phase_deg = -[57.19  58.19  61.39  59.09  59.89  51.19  46.89  42.79  36.89  32.0  44.0  32.0  37.59  45.29  50.09][:]
σ_log10_ρ = [0.0721  0.0425  0.0244  0.021  0.0164  0.0173  0.0287  0.0328  0.0193  0.027  0.0591  0.0506  0.0825  0.1233  0.0927][:]
σ_phase_deg = [22.95  22.95  4.96  4.46  5.96  22.95  22.95  2.46  1.65  22.95  6.37  2.46  22.95  4.15  22.95][:]
F = transD_GP.MT1DInversion.MT1DZ_nodepthprior(d_log10_ρ = d_log10_ρ, d_phase_deg = d_phase_deg,
                                    σ_log10_ρ = σ_log10_ρ, σ_phase_deg = σ_phase_deg,
                                    freqs = 1 ./T, zboundaries = zboundaries, irxlayer = 1)
transD_GP.MT1DInversion.plotdata(F)


function transit_orb(t,x,dt,nsub)
# Computes a transit lightcurve, normalized to unity, as a
# function of time, t, usually in HJD or BJD.
#
# Input parameters (x) are:
# x[1] = P  (units of day)
# x[2] = inc = inclination angle
# x[3] = p = R_p/R_* = radius of planet in units of radius of star
# x[4] = t0 = mid-point of transit
# x[5] = u1 = linear limb-darkening coefficient
# x[6] = u2 = quadratic limb-darkening coefficient
# x[7] = f0 = uneclipsed flux
# x[8] = a/R_* = semi-major axis divided by R_*
# x[9] = e = eccentricity
# x[10] = omega = longitude of pericentre
# Compute time of pericentre passage:
# The true anomaly at the time of transit:
f1=1.5*pi-x[10]*pi/180.0
ecc=x[9]
tp=(x[4]+x[1]*sqrt(1.0-ecc^2)/2.0/pi*(ecc*sin(f1)/(1.0+ecc*cos(f1))
    -2.0/sqrt(1.0-ecc^2)*atan2(sqrt(1.0-ecc^2)*tan(0.5*f1),1.0+ecc)))
fluxoft = 0.0
for j=1:nsub
  tsub = t-dt/2.0+dt*(j-0.5)/nsub
  m=2.0*pi/x[1]*(tsub-tp)
  f=kepler(m,ecc)
  radius=x[8]*(1.0-ecc^2)/(1.0+ecc*cos(f))
# Now compute sky separation:
  z0=radius*sqrt(1.0-(sin(x[2]*pi/180.0)*sin(x[10]*pi/180.0+f))^2)
  if (sin(x[10]*pi/180.0+f) < 0.0) && (z0 <= (1.0+x[3]))
    if x[3] < 1.0
      fluxoft += occultquad(z0,x[5],x[6],x[3])
    else
# We'll assume that the smaller object is not limb-darkened:
      fluxoft += occultquad(z0,  0.,  0.,x[3])
    end
  else
    fluxoft += 1.0
  end
end
# Take mean of flux over sub-time steps:
flux = fluxoft*x[7]/nsub
return flux
end


module BoundaryValueDiffEq

using Reexport, LinearAlgebra, SparseArrays
@reexport using DiffEqBase

import DiffEqBase: solve
import NLsolve, ForwardDiff, BandedMatrices, FiniteDiff

struct MIRKTableau{T}
    c::Vector{T}
    v::Vector{T}
    b::Vector{T}
    x::Matrix{T}
end

# ODE BVP problem system
mutable struct BVPSystem{T,U<:AbstractArray,P}
    order::Int                  # The order of MIRK method
    M::Int                      # Number of equations in the ODE system
    N::Int                      # Number of nodes in the mesh
    fun!                        # M -> M
    bc!                         # 2 -> 2
    p::P
    x::Vector{T}                # N
    y::Vector{U}                # N{M}
    f::Vector{U}                # N{M}
    residual::Vector{U}         # N{M}
    tmp::Vector{T}
end

include("vector_auxiliary.jl")
include("algorithms.jl")
include("alg_utils.jl")
include("mirk_tableaus.jl")
include("cache.jl")
include("collocation.jl")
include("jacobian.jl")
include("solve.jl")

export Shooting
export MIRK4, GeneralMIRK4

end # module


struct JC69rel <: JC69 end

function show(io::IO, object::JC69rel)
  print(io, "\r\e[0m\e[1mJ\e[0mukes and \e[1mC\e[0mantor 19\e[1m69\e[0m model (relative rate form)")
end

JC69() = JC69rel()

@inline function Q(mod::JC69rel)
  Q₁ =  0.25
  Q₂ = -(Q₁ * 3)
  return SMatrix{4, 4, Float64}(Q₂, Q₁, Q₁, Q₁,
                                Q₁, Q₂, Q₁, Q₁,
                                Q₁, Q₁, Q₂, Q₁,
                                Q₁, Q₁, Q₁, Q₂)
end

@inline function P(mod::JC69rel, t::Float64)
  if t < 0.0
    error("t must be positive")
  end
  e₁ = exp(-t)
  P₁ = 0.25 + 0.75 * e₁
  P₂ = 0.25 - 0.25 * e₁
  return Pmatrix(P₁, P₂, P₂, P₂,
                 P₂, P₁, P₂, P₂,
                 P₂, P₂, P₁, P₂,
                 P₂, P₂, P₂, P₁)
end


function velocityphi(r,x,b,mode,d)

#   evaluates the velocity from the solution
#
#   input
#   r        - radial grid (excludes boundaries)
#   x        - cos(theta) grid (excludes poles and equator)
#   b        - radial magnetic field (mT)
#   eta      - dimensionless diffusivity
#   mode     - solution vector (eigenfunction)
#   d        - eigenvalue  (frequency = d / 1im)

#   some constants
    omega = 0.7292e-4;       # rotation rate (1/s)
    sectoyear = 3.155e7;     # seconds in a year
    mu = 4 * pi * 1.0e-7;    # permeability
    eta = 0.8;            # magnetic diffusivity  (m^2/s)
    rho = 1.0e4;             # density
    L = 3.480e6;             # core radius (length scale)
    bscale = sqrt(rho * mu) * omega * L;

#   dimensionless quantities
    B = b * 1.0e-3 / bscale;
    Eta = eta / (omega * L^2);

#  allocate memory for velocity field
   nr = length(r);
   nx = length(x);
   ndim = nr * nx;
   dvdr = Array{ComplexF64}(undef,ndim);
   vphi = Array{ComplexF64}(undef,ndim);

#  radial magnetic field (dimensionless)
   br = magnetic_field(B,x);

#  retrieve required matrices
   dsp = diffusion(r,x,eta);

#  rate of change of bphi
   dbdt = d * mode[1:ndim];

#  account for diffusion
   dbdr2 = dsp[ndim+1:2*ndim,ndim+1:2*ndim] * mode[1:ndim];

#  combine to evaluate velocity gradient
   f = dbdt - dbdr2;
   for l = 1 : nx
       for k = 1 : nr
           ipos = (l-1) * nr + k;
           dvdr[ipos] = f[ipos] / br[l];
       end
   end

#  integrate for vphi
   dr = r[2] - r[1];
   for l = 1 : nx
       for k = 2 : nr
           ipos = (l-1) * nr + k;
           vphi[ipos] = vphi[ipos-1]+ 0.5*(dvdr[ipos] + dvdr[ipos-1])*dr;
       end
   end

#  extract CMB value
   vphicmb = Array{ComplexF64}(undef,nx)
   bphicmb = Array{ComplexF64}(undef,nx);

   for l = 1 : nx
       k = nr;
       ipos = (l-1)*nr + k;
#       println(" l ", l, " ipos ",ipos," vphi ",vphi[ipos]," ",vphi[ipos-1])
       vphicmb[l] = vphi[ipos];
       bphicmb[l] = mode[ipos];
#       println(" l ", vphicmb[l])
   end

   return vphicmb,bphicmb

end


numTests = 5
maxDepth = 50
numColors = 6
target = [100.0, 100.0, 100.0]
fnEvals=0
# alphaScoreMatrix = 10 .* rand(numColors+1,numColors+1) .- 5
# betaScoreMatrix = 10 .* rand(numColors+1,numColors+1) .- 5
# gammaScoreMatrix = 10 .* rand(numColors+1,numColors+1) .- 5

using LinearAlgebra
using Plotly
include("EvaluationFunctions.jl")
include("GenerationFunctions.jl")
#generate = GenerateRandomly
generate = ExhaustiveHillClimbing

scores = Array{Float64, 2}(undef, numTests, maxDepth)
for i=1:numTests

    global    alphaScoreMatrix = 10 .* rand(numColors+1,numColors+1) .- 5
    global    betaScoreMatrix = 10 .* rand(numColors+1,numColors+1) .- 5
    global    gammaScoreMatrix = 10 .* rand(numColors+1,numColors+1) .- 5
    testTree1 = generate(1)
    scores[i,1]= GetCombinedScore(testTree1)
    for j=2:maxDepth
        testTree1 = generate(1, cand=testTree1)
        scores[i,j]= GetCombinedScore(testTree1)
    end
end

plot(transpose(scores))


# eliminate children
function collapse(ht::HTensorNode{T}) where {T}
   if isleaf(ht); return ht; end
   core = coreshape(ht)
   nchildren = length(ht.children)
   newsubshape = Int[]
   newchildren = HTensorNode{T}[]
   for i in 1:nchildren
      child = ht.children[i]
      core = tensordot(core, child.data, i)
      append!(newsubshape, child.subshape)
      append!(newchildren, child.children)
   end
  #println(STDERR, "!!! collapse: $(ht.subshape) -> $(newsubshape)")
   core = reshape(core, newsubshape..., ht.nbasis)
   HTensorNode(core, newchildren)
end

# eliminate children recursively
function collapse_all(ht::HTensorNode)
   if isleaf(ht); return ht; end
   newchildren = map(collapse_all, ht.children)
   collapse(HTensorNode(ht.data, newchildren))
end


using DelimitedFiles
using Test

"""
    count_trees(data, move=(1, 3))

Move and count trees according to part 1's rules.
"""
function count_trees(data, move=(1, 3))::Int
    arr = permutedims(hcat(collect.(data)...))  # wtf syntax
    rows, cols = size(arr)
    dr, dc = move
    total = 0
    col = 1
    for row in 1:dr:rows
        if arr[row, col] == '#'
            total += 1
            end
        col = (col + dc - 1) % cols + 1  # Stupid 1-index, next time: mod1(a, b)
        end
    return total
    end

test = readdlm("./data/day03_test.txt")
@test count_trees(test) == 7

data = readdlm("./data/day03.txt")
println("Part 1 trees: $(count_trees(data))")

"""
    count_all(data, moves)

Part 2 function
"""
function count_all(data, moves)
    product = 1
    for move in moves
        product *= count_trees(data, move)
        end
    return product
    end

moves = [(1, 1), (1, 3), (1, 5), (1, 7), (2, 1)]

@test count_all(test, moves) == 336

println("Part 2 product: $(count_all(data, moves))")


"""
Copy of Cosmology.jl to remove Unitful dependence.

Original license:

The MIT License (MIT)

Copyright (c) 2013 Mike Nolta <mike@nolta.net>

Permission is hereby granted, free of charge, to any person obtaining a
copy of this software and associated documentation files (the "Software"),
to deal in the Software without restriction, including without limitation
the rights to use, copy, modify, merge, publish, distribute, sublicense,
and/or sell copies of the Software, and to permit persons to whom the
Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
DEALINGS IN THE SOFTWARE.
"""
module CosmologyTools

using Slomo.Constants: rho_crit_h2
export Ωm, ρcrit, ρm, default_cosmo

abstract type AbstractCosmology end
abstract type AbstractClosedCosmology <: AbstractCosmology end
abstract type AbstractFlatCosmology <: AbstractCosmology end
abstract type AbstractOpenCosmology <: AbstractCosmology end

struct FlatLCDM{T <: Real} <: AbstractFlatCosmology
    h::T
    Ω_Λ::T
    Ω_m::T
    Ω_r::T
end
FlatLCDM(h::Real, Ω_Λ::Real, Ω_m::Real, Ω_r::Real) =
    FlatLCDM(promote(float(h), float(Ω_Λ), float(Ω_m), float(Ω_r))...)

a2E(c::FlatLCDM, a) = sqrt(c.Ω_r + c.Ω_m * a + c.Ω_Λ * a^4)

struct ClosedLCDM{T <: Real} <: AbstractClosedCosmology
    h::T
    Ω_k::T
    Ω_Λ::T
    Ω_m::T
    Ω_r::T
end
ClosedLCDM(h::Real, Ω_k::Real, Ω_Λ::Real, Ω_m::Real, Ω_r::Real) =
    ClosedLCDM(promote(float(h), float(Ω_k), float(Ω_Λ), float(Ω_m),
                       float(Ω_r))...)

struct OpenLCDM{T <: Real} <: AbstractOpenCosmology
    h::T
    Ω_k::T
    Ω_Λ::T
    Ω_m::T
    Ω_r::T
end
OpenLCDM(h::Real, Ω_k::Real, Ω_Λ::Real, Ω_m::Real, Ω_r::Real) =
    OpenLCDM(promote(float(h), float(Ω_k), float(Ω_Λ), float(Ω_m),
                     float(Ω_r))...)

function a2E(c::Union{ClosedLCDM,OpenLCDM}, a)
    a2 = a * a
    sqrt(c.Ω_r + c.Ω_m * a + (c.Ω_k + c.Ω_Λ * a2) * a2)
end

for c in ("Flat", "Open", "Closed")
    name = Symbol("$(c)WCDM")
    @eval begin
        struct $(name){T <: Real} <: $(Symbol("Abstract$(c)Cosmology"))
            h::T
            Ω_k::T
            Ω_Λ::T
            Ω_m::T
            Ω_r::T
            w0::T
            wa::T
        end
        function $(name)(h::Real, Ω_k::Real, Ω_Λ::Real, Ω_m::Real, Ω_r::Real,
                         w0::Real, wa::Real)
            $(name)(promote(float(h), float(Ω_k), float(Ω_Λ), float(Ω_m),
                            float(Ω_r), float(w0), float(wa))...)
        end
    end
end

function WCDM(h::Real, Ω_k::Real, Ω_Λ::Real, Ω_m::Real, Ω_r::Real, w0::Real, wa::Real)
    if Ω_k < 0
        ClosedWCDM(h, Ω_k, Ω_Λ, Ω_m, Ω_r, w0, wa)
    elseif Ω_k > 0
        OpenWCDM(h, Ω_k, Ω_Λ, Ω_m, Ω_r, w0, wa)
    else
        FlatWCDM(h, Ω_k, Ω_Λ, Ω_m, Ω_r, w0, wa)
    end
end

function a2E(c::Union{FlatWCDM,ClosedWCDM,OpenWCDM}, a)
    ade = exp((1 - 3 * (c.w0 + c.wa)) * log(a) + 3 * c.wa * (a - 1))
    sqrt(c.Ω_r + (c.Ω_m + c.Ω_k * a) * a + c.Ω_Λ * ade)
end

"""
    cosmology(;h = 0.69,
               Neff = 3.04,
               OmegaK = 0,
               OmegaM = 0.29,
               OmegaR = nothing,
               Tcmb = 2.7255,
               w0 = -1,
               wa = 0)
# Parameters
* `h` - Dimensionless Hubble constant
* `OmegaK` - Curvature density (Ω_k)
* `OmegaM` - Matter density (Ω_m)
* `OmegaR` - Radiation density (Ω_r)
* `Tcmb` - CMB temperature in Kelvin; used to compute Ω_γ
* `Neff` - Effective number of massless neutrino species; used to compute Ω_ν
* `w0` - CPL dark energy equation of state; `w = w0 + wa(1-a)`
* `wa` - CPL dark energy equation of state; `w = w0 + wa(1-a)`
# Examples
```jldoctest
julia> c = cosmology()
Cosmology.FlatLCDM{Float64}(0.69, 0.7099122024007928, 0.29, 8.77975992071536e-5)
julia> c = cosmology(OmegaK=0.1)
Cosmology.OpenLCDM{Float64}(0.69, 0.1, 0.6099122024007929, 0.29, 8.77975992071536e-5)
julia> c = cosmology(w0=-0.9, OmegaK=-0.1)
Cosmology.ClosedWCDM{Float64}(0.69, -0.1, 0.8099122024007929, 0.29, 8.77975992071536e-5, -0.9, 0.0)
```
"""
function cosmology(;h = 0.69,
                   Neff = 3.04,
                   OmegaK = 0,
                   OmegaM = 0.29,
                   OmegaR = nothing,
                   Tcmb = 2.7255,
                   w0 = -1,
                   wa = 0)

    if OmegaR === nothing
        OmegaG = 4.48131e-7 * Tcmb^4 / h^2
        OmegaN = Neff * OmegaG * (7 / 8) * (4 / 11)^(4 / 3)
        OmegaR = OmegaG + OmegaN
    end

    OmegaL = 1 - OmegaK - OmegaM - OmegaR

    if !(w0 == -1 && wa == 0)
        return WCDM(h, OmegaK, OmegaL, OmegaM, OmegaR, w0, wa)
    end

    if OmegaK < 0
        return ClosedLCDM(h, OmegaK, OmegaL, OmegaM, OmegaR)
        elseif OmegaK > 0
        return OpenLCDM(h, OmegaK, OmegaL, OmegaM, OmegaR)
        else
        return FlatLCDM(h, OmegaL, OmegaM, OmegaR)
    end
end

# hubble rate
scale_factor(z) = 1 / (1 + z)
E(c::AbstractCosmology, z) = (a = scale_factor(z); a2E(c, a) / a^2)

const planck18 = cosmology(h = 0.6766,
                           Neff = 3.046,
                           OmegaM = 0.3111)

const planck15 = cosmology(h = 0.6774,
                           Neff = 3.046,
                           OmegaM = 0.3089)

const planck13 = cosmology(h = 0.6777,
                           Neff = 3.046,
                           OmegaM = 0.3071)

const wmap9 = cosmology(h = 0.6932,
                        Neff = 3.046,
                        OmegaM = 0.2865)

default_cosmo = planck18

"""
Compute the matter density of the universe at redshift z, in units of the critical density.
"""
function Ωm(z; cosmo = default_cosmo)
    return cosmo.Ω_m * (1.0 + z) ^ 3 / E(cosmo, z) ^ 2
end

"""
Compute the critical density of the universe at redshift z, in Msun / kpc3
"""
function ρcrit(z; cosmo = default_cosmo)
    return rho_crit_h2 * cosmo.h^2 * E(cosmo, z) ^ 2
end

"""
Compute the mass density of the universe at redshift z, in Msun / kpc3
ρm(z) = Ωm(z) * ρcrit(z)
"""
function ρm(z; cosmo = default_cosmo)
    return cosmo.Ω_m * (1 + z)^3 * rho_crit_h2 * cosmo.h^2
end

end


# https://adventofcode.com/2020/day/14

import Base: ==

example1 = "
mask = XXXXXXXXXXXXXXXXXXXXXXXXXXXXX1XXXX0X
mem[8] = 11
mem[7] = 101
mem[8] = 0
"

const MEM_SIZE = 36

BitMask = Vector{Union{Nothing,Bool}}
Binary = Vector{Union{Bool}}

Base.repr(obj::BitMask) = join([Dict(nothing => 'X', true => 1, false => 0)[x] for x in obj])

struct Program
    mask::BitMask
    mem::Vector{Tuple{Integer,Integer}}
end

(==)(x::Program, y::Program) = (x.mask == y.mask) && (x.mem == y.mem)

function parse_input(string)
    programs = Program[]
    mask, mem = [], []
    for (i, row) in enumerate(split(string, '\n', keepempty=false))
        row = strip(row)
        key, value = split(row, '=')
        key, value = strip(key), strip(value)

        if key == "mask"
            if i > 1
                push!(programs, Program(mask, mem))
            end

            mask = map(x -> x == 'X' ? nothing : parse(Bool, x), collect(value))
            @assert length(mask) == MEM_SIZE
            mem = []
        else
            address = parse(Int, match(r"mem\[(\d+)\]", key).captures[1])
            value = parse(Int, value)
            push!(mem, (address, value))
        end
    end

    push!(programs, Program(mask, mem))
    return programs
end

@assert parse_input(example1) ==
        [Program([repeat([nothing], 29); 1; repeat([nothing], 4); 0; nothing], [(8, 11), (7, 101), (8, 0)])]
@assert repr(parse_input(example1)[1].mask) == "XXXXXXXXXXXXXXXXXXXXXXXXXXXXX1XXXX0X"

test_program = "
mask = XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
mem[0] = 0
mask = XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
mem[1] = 1
mem[2] = 2
"

@assert parse_input(test_program) ==
        [Program(repeat([nothing], MEM_SIZE), [(0, 0)]), Program(repeat([nothing], MEM_SIZE), [(1, 1), (2, 2)])]

function binary_to_decimal(binary::Binary)
    decimal = 0
    for (i, x) in enumerate(reverse(binary))
        decimal += 2^(i - 1) * x
    end
    return decimal
end

@assert binary_to_decimal(Bool[0, 0, 0]) == 0
@assert binary_to_decimal(Bool[0, 0, 1]) == 1
@assert binary_to_decimal(Bool[0, 1, 0]) == 2
@assert binary_to_decimal(Bool[1, 0, 0, 0, 0, 0, 0]) == 64
@assert binary_to_decimal(Bool[1, 1, 0, 0, 1, 0, 1]) == 101

function apply_mask(bits::Binary, mask::BitMask)::Binary
    idx = map(x -> !isnothing(x), mask)
    bits[idx] = mask[idx]
    return bits
end

@assert apply_mask(Bool[0, 1, 0, 1], [nothing, nothing, nothing, true]) == Bool[0, 1, 0, 1]
@assert apply_mask(Bool[0, 0, 0, 0], [nothing, true, nothing, false]) == Bool[0, 1, 0, 0]
@assert apply_mask(Bool[0, 0, 0, 0], [true, true, true, nothing]) == Bool[1, 1, 1, 0]
@assert apply_mask(Bool[0, 1, 0, 1], [true, nothing, true, nothing]) == Bool[1, 1, 1, 1]

function tobinary(value::Integer)::Binary
    return convert(Vector{Bool}, reverse(digits(value, base=2, pad=MEM_SIZE)))
end

@assert tobinary(42) ==
        Bool[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 0, 1, 0]

function write(programs::Vector{Program})
    memory = Dict{Integer,Integer}()
    for program in programs
        for cmd in program.mem
            bits = tobinary(cmd[2])
            bits = apply_mask(bits, program.mask)
            memory[cmd[1]] = binary_to_decimal(bits)
        end
    end
    return memory
end

@assert write(parse_input(example1)) == Dict(7 => 101, 8 => 64)

function part1(input)
    program = parse_input(input)
    memory = write(program)
    sum(values(memory))
end

@assert part1(example1) == 165

"""
Apply mask:

```
address: 000000000000000000000000000000101010  (decimal 42)
mask:    000000000000000000000000000000X1001X
result:  000000000000000000000000000000X1101X
```
"""
function apply_xmask(bits::Binary, mask::BitMask)::BitMask
    @assert length(bits) == length(mask)
    bits = convert(BitMask, bits)
    bits[mask .== true] .= true
    bits[isnothing.(mask)] .= nothing
    return bits
end

@assert apply_xmask(tobinary(42), [repeat([false], 30); nothing; true; false; false; true; nothing]) ==
        [repeat([false], 30); nothing; true; true; false; true; nothing]

Base.repr(obj::Binary) = join([Dict(true => 1, false => 0)[x] for x in obj])

function write_ver2(programs::Vector{Program}; verbose=false)
    memory = Dict{Integer,Integer}()
    for program in programs
        for cmd in program.mem
            bits = tobinary(cmd[1])

            verbose && println("address: $(repr(bits)) (decimal $(binary_to_decimal(bits)))")
            verbose && println("mask:    $(repr(program.mask))")

            bits = apply_xmask(bits, program.mask)
            verbose && println("result:  $(repr(bits))\n")

            if any(isnothing.(bits))
                idx = findall(isnothing, bits)
                for replacement in Iterators.product(repeat([0:1], length(idx))...)
                    bits[idx] .= convert(Vector{Bool}, collect(replacement))
                    address = binary_to_decimal(convert(Binary, bits))
                    memory[address] = cmd[2]
                    verbose && println("$(join(convert(Vector{Int}, bits))) (decimal $(address))")
                end
            else
                address = binary_to_decimal(bits)
                memory[address] = cmd[2]
                verbose && println("$(join(convert(Vector{Int}, bits))) (decimal $(address))")
            end
        end
        verbose && println()
    end
    return memory
end

example2 = "
mask = 000000000000000000000000000000X1001X
mem[42] = 100
mask = 00000000000000000000000000000000X0XX
mem[26] = 1
"

function part2(input)
    program = parse_input(input)
    memory = write_ver2(program)
    sum(values(memory))
end

@assert part2(example2) == 208

test = read("data/day-14.txt", String)
println("Part 1: $(result1 = part1(test))")
println("Part 2: $(result2 = part2(test))")

@assert result1 == 6386593869035
@assert result2 == 4288986482164


using MultipleViewGeometry, Test
using MultipleViewGeometry.ModuleTypes

ℳ = [Point2DH(rand(1,3)) for i = 1:10000]

moments(FundamentalMatrix(), (ℳ,ℳ))

# TODO


@inline function initialize!(integrator,cache::BS3ConstantCache,f=integrator.f)
  integrator.kshortsize = 2
  integrator.k = eltype(integrator.sol.k)(integrator.kshortsize)
  integrator.fsalfirst = f(integrator.t,integrator.uprev) # Pre-start fsal
end

@inline function perform_step!(integrator,cache::BS3ConstantCache,f=integrator.f)
  @unpack t,dt,uprev,u,k = integrator
  @unpack a21,a32,a41,a42,a43,c1,c2,b1,b2,b3,b4 = cache
  k1 = integrator.fsalfirst
  a1 = dt*a21
  k2 = f(@muladd(t+c1*dt),@muladd(uprev+a1*k1))
  a2 = dt*a32
  k3 = f(@muladd(t+c2*dt),@muladd(uprev+a2*k2))
  u = @muladd uprev+dt*(a41*k1+a42*k2+a43*k3)
  k4 = f(t+dt,u); integrator.fsallast = k4
  if integrator.opts.adaptive
    utilde = @muladd uprev + dt*(b1*k1 + b2*k2 + b3*k3 + b4*k4)
    integrator.EEst = integrator.opts.internalnorm( ((utilde-u)./@muladd(integrator.opts.abstol+max(abs.(uprev),abs.(u))*integrator.opts.reltol)))
  end
  integrator.k[1] = integrator.fsalfirst
  integrator.k[2] = integrator.fsallast
  @pack integrator = t,dt,u
end

@inline function initialize!(integrator,cache::BS3Cache,f=integrator.f)
  integrator.kshortsize = 2
  integrator.k = eltype(integrator.sol.k)(integrator.kshortsize)
  integrator.fsalfirst = cache.fsalfirst  # done by pointers, no copying
  integrator.fsallast = cache.k4
  integrator.k[1] = integrator.fsalfirst
  integrator.k[2] = integrator.fsallast
  f(integrator.t,integrator.uprev,integrator.fsalfirst) # Pre-start fsal
end

@inline function perform_step!(integrator,cache::BS3Cache,f=integrator.f)
  @unpack t,dt,uprev,u,k = integrator
  uidx = eachindex(integrator.uprev)
  @unpack k2,k3,k4,utilde,tmp,atmp = cache
  k1 = cache.fsalfirst
  @unpack a21,a32,a41,a42,a43,c1,c2,b1,b2,b3,b4 = cache.tab
  a1 = dt*a21
  for i in uidx
    tmp[i] = @muladd uprev[i]+a1*k1[i]
  end
  f(@muladd(t+c1*dt),tmp,k2)
  a2 = dt*a32
  for i in uidx
    tmp[i] = @muladd uprev[i]+a2*k2[i]
  end
  f(@muladd(t+c2*dt),tmp,k3)
  for i in uidx
    u[i] = @muladd uprev[i]+dt*(a41*k1[i]+a42*k2[i]+a43*k3[i])
  end
  f(t+dt,u,k4)
  if integrator.opts.adaptive
    for i in uidx
      utilde[i] = @muladd uprev[i] + dt*(b1*k1[i] + b2*k2[i] + b3*k3[i] + b4*k4[i])
      atmp[i] = ((utilde[i]-u[i])./@muladd(integrator.opts.abstol+max(abs(uprev[i]),abs(u[i]))*integrator.opts.reltol))
    end
    integrator.EEst = integrator.opts.internalnorm(atmp)
  end
  @pack integrator = t,dt,u,k
end

@inline function initialize!(integrator,cache::BS5ConstantCache,f=integrator.f)
  integrator.kshortsize = 8
  integrator.k = eltype(integrator.sol.k)(integrator.kshortsize)
  integrator.fsalfirst = f(integrator.t,integrator.uprev) # Pre-start fsal
end

@inline function perform_step!(integrator,cache::BS5ConstantCache,f=integrator.f)
  @unpack t,dt,uprev,u,k = integrator
        @unpack c1,c2,c3,c4,c5,a21,a31,a32,a41,a42,a43,a51,a52,a53,a54,a61,a62,a63,a64,a65,a71,a72,a73,a74,a75,a76,a81,a83,a84,a85,a86,a87,bhat1,bhat3,bhat4,bhat5,bhat6,btilde1,btilde2,btilde3,btilde4,btilde5,btilde6,btilde7,btilde8 = cache
  k1 = integrator.fsalfirst
  a = dt*a21
  k2 = f(@muladd(t+c1*dt),@muladd(uprev+a*k1))
  k3 = f(@muladd(t+c2*dt),@muladd(uprev+dt*(a31*k1+a32*k2)))
  k4 = f(@muladd(t+c3*dt),@muladd(uprev+dt*(a41*k1+a42*k2+a43*k3)))
  k5 = f(@muladd(t+c4*dt),@muladd(uprev+dt*(a51*k1+a52*k2+a53*k3+a54*k4)))
  k6 = f(@muladd(t+c5*dt),@muladd(uprev+dt*(a61*k1+a62*k2+a63*k3+a64*k4+a65*k5)))
  k7 = f(t+dt,@muladd(uprev+dt*(a71*k1+a72*k2+a73*k3+a74*k4+a75*k5+a76*k6)))
  u = @muladd uprev+dt*(a81*k1+a83*k3+a84*k4+a85*k5+a86*k6+a87*k7)
  integrator.fsallast = f(t+dt,u); k8 = integrator.fsallast
  if integrator.opts.adaptive
    uhat   = dt*@muladd(bhat1*k1 + bhat3*k3 + bhat4*k4 + bhat5*k5 + bhat6*k6)
    utilde = @muladd uprev + dt*(btilde1*k1 + btilde2*k2 + btilde3*k3 + btilde4*k4 + btilde5*k5 + btilde6*k6 + btilde7*k7 + btilde8*k8)
    EEst1 = integrator.opts.internalnorm( sum(((uhat)./@muladd(integrator.opts.abstol+max(abs.(uprev),abs.(u))*integrator.opts.reltol))))
    EEst2 = integrator.opts.internalnorm( sum(((utilde-u)./@muladd(integrator.opts.abstol+max(abs.(uprev),abs.(u))*integrator.opts.reltol))))
    integrator.EEst = max(EEst1,EEst2)
  end
  integrator.k[1]=k1; integrator.k[2]=k2; integrator.k[3]=k3;integrator.k[4]=k4;integrator.k[5]=k5;integrator.k[6]=k6;integrator.k[7]=k7;integrator.k[8]=k8
  @pack integrator = t,dt,u,k
end

@inline function initialize!(integrator,cache::BS5Cache,f=integrator.f)
  integrator.kshortsize = 8
  integrator.k = eltype(integrator.sol.k)(integrator.kshortsize)
  integrator.k = eltype(integrator.sol.k)(integrator.kshortsize)
  integrator.k[1]=cache.k1; integrator.k[2]=cache.k2;
  integrator.k[3]=cache.k3; integrator.k[4]=cache.k4;
  integrator.k[5]=cache.k5; integrator.k[6]=cache.k6;
  integrator.k[7]=cache.k7; integrator.k[8]=cache.k8
  integrator.fsalfirst = cache.k1; integrator.fsallast = cache.k8  # setup pointers
  f(integrator.t,integrator.uprev,integrator.fsalfirst) # Pre-start fsal
end

@inline function perform_step!(integrator,cache::BS5Cache,f=integrator.f)
  @unpack t,dt,uprev,u,k = integrator
  uidx = eachindex(integrator.uprev)
  @unpack k1,k2,k3,k4,k5,k6,k7,k8,utilde,uhat,tmp,atmp,atmptilde = cache
  @unpack c1,c2,c3,c4,c5,a21,a31,a32,a41,a42,a43,a51,a52,a53,a54,a61,a62,a63,a64,a65,a71,a72,a73,a74,a75,a76,a81,a83,a84,a85,a86,a87,bhat1,bhat3,bhat4,bhat5,bhat6,btilde1,btilde2,btilde3,btilde4,btilde5,btilde6,btilde7,btilde8 = cache.tab
  a = dt*a21
  for i in uidx
    tmp[i] = @muladd uprev[i]+a*k1[i]
  end
  f(@muladd(t+c1*dt),tmp,k2)
  for i in uidx
    tmp[i] = @muladd uprev[i]+dt*(a31*k1[i]+a32*k2[i])
  end
  f(@muladd(t+c2*dt),tmp,k3)
  for i in uidx
    tmp[i] = @muladd uprev[i]+dt*(a41*k1[i]+a42*k2[i]+a43*k3[i])
  end
  f(@muladd(t+c3*dt),tmp,k4)
  for i in uidx
    tmp[i] = @muladd uprev[i]+dt*(a51*k1[i]+a52*k2[i]+a53*k3[i]+a54*k4[i])
  end
  f(@muladd(t+c4*dt),tmp,k5)
  for i in uidx
    tmp[i] = @muladd uprev[i]+dt*(a61*k1[i]+a62*k2[i]+a63*k3[i]+a64*k4[i]+a65*k5[i])
  end
  f(@muladd(t+c5*dt),tmp,k6)
  for i in uidx
    tmp[i] = @muladd uprev[i]+dt*(a71*k1[i]+a72*k2[i]+a73*k3[i]+a74*k4[i]+a75*k5[i]+a76*k6[i])
  end
  f(t+dt,tmp,k7)
  for i in uidx
    u[i] = @muladd uprev[i]+dt*(a81*k1[i]+a83*k3[i]+a84*k4[i]+a85*k5[i]+a86*k6[i]+a87*k7[i])
  end
  f(t+dt,u,k8)
  if integrator.opts.adaptive
    for i in uidx
      uhat[i]   = dt*@muladd(bhat1*k1[i] + bhat3*k3[i] + bhat4*k4[i] + bhat5*k5[i] + bhat6*k6[i])
      utilde[i] = @muladd uprev[i] + dt*(btilde1*k1[i] + btilde2*k2[i] + btilde3*k3[i] + btilde4*k4[i] + btilde5*k5[i] + btilde6*k6[i] + btilde7*k7[i] + btilde8*k8[i])
      atmp[i] = ((uhat[i])./@muladd(integrator.opts.abstol+max(abs(uprev[i]),abs(u[i]))*integrator.opts.reltol))
      atmptilde[i] = ((utilde[i]-u[i])./@muladd(integrator.opts.abstol+max(abs(uprev[i]),abs(u[i]))*integrator.opts.reltol))
    end
    EEst1 = integrator.opts.internalnorm(atmp)
    EEst2 = integrator.opts.internalnorm(atmptilde)
    integrator.EEst = max(EEst1,EEst2)
  end
  @pack integrator = t,dt,u,k
end

@inline function initialize!(integrator,cache::Tsit5ConstantCache,f=integrator.f)
  integrator.kshortsize = 7
  integrator.k = eltype(integrator.sol.k)(integrator.kshortsize)
  integrator.fsalfirst = f(integrator.t,integrator.uprev) # Pre-start fsal
end

@inline function perform_step!(integrator,cache::Tsit5ConstantCache,f=integrator.f)
  @unpack t,dt,uprev,u,k = integrator
  @unpack c1,c2,c3,c4,c5,c6,a21,a31,a32,a41,a42,a43,a51,a52,a53,a54,a61,a62,a63,a64,a65,a71,a72,a73,a74,a75,a76,b1,b2,b3,b4,b5,b6,b7 = cache
  k1 = integrator.fsalfirst
  a = dt*a21
  k2 = f(@muladd(t+c1*dt),@muladd(uprev+a*k1))
  k3 = f(@muladd(t+c2*dt),@muladd(uprev+dt*(a31*k1+a32*k2)))
  k4 = f(@muladd(t+c3*dt),@muladd(uprev+dt*(a41*k1+a42*k2+a43*k3)))
  k5 = f(@muladd(t+c4*dt),@muladd(uprev+dt*(a51*k1+a52*k2+a53*k3+a54*k4)))
  k6 = f(t+dt,@muladd(uprev+dt*(a61*k1+a62*k2+a63*k3+a64*k4+a65*k5)))
  u = @muladd uprev+dt*(a71*k1+a72*k2+a73*k3+a74*k4+a75*k5+a76*k6)
  integrator.fsallast = f(t+dt,u); k7 = integrator.fsallast
  if integrator.opts.adaptive
    utilde = @muladd uprev + dt*(b1*k1 + b2*k2 + b3*k3 + b4*k4 + b5*k5 + b6*k6 + b7*k7)
    integrator.EEst = integrator.opts.internalnorm(((utilde-u)./@muladd(integrator.opts.abstol+max(abs.(uprev),abs.(u))*integrator.opts.reltol)))
  end
  integrator.k[1] = k1
  integrator.k[2] = k2
  integrator.k[3] = k3
  integrator.k[4] = k4
  integrator.k[5] = k5
  integrator.k[6] = k6
  integrator.k[7] = k7
  @pack integrator = t,dt,u,k
end

@inline function initialize!(integrator,cache::Tsit5Cache,f=integrator.f)
  integrator.kshortsize = 7
  integrator.fsalfirst = cache.k1; integrator.fsallast = cache.k7 # setup pointers
  integrator.k = eltype(integrator.sol.k)(integrator.kshortsize)
  # Setup k pointers
  integrator.k[1] = cache.k1
  integrator.k[2] = cache.k2
  integrator.k[3] = cache.k3
  integrator.k[4] = cache.k4
  integrator.k[5] = cache.k5
  integrator.k[6] = cache.k6
  integrator.k[7] = cache.k7
  f(integrator.t,integrator.uprev,integrator.fsalfirst) # Pre-start fsal
end

@inline function perform_step!(integrator,cache::Tsit5Cache,f=integrator.f)
  @unpack t,dt,uprev,u,k = integrator
  uidx = eachindex(integrator.uprev)
  @unpack c1,c2,c3,c4,c5,c6,a21,a31,a32,a41,a42,a43,a51,a52,a53,a54,a61,a62,a63,a64,a65,a71,a72,a73,a74,a75,a76,b1,b2,b3,b4,b5,b6,b7 = cache.tab
  @unpack k1,k2,k3,k4,k5,k6,k7,utilde,tmp,atmp = cache
  a = dt*a21
  for i in uidx
    tmp[i] = @muladd uprev[i]+a*k1[i]
  end
  f(@muladd(t+c1*dt),tmp,k2)
  for i in uidx
    tmp[i] = @muladd uprev[i]+dt*(a31*k1[i]+a32*k2[i])
  end
  f(@muladd(t+c2*dt),tmp,k3)
  for i in uidx
    tmp[i] = @muladd uprev[i]+dt*(a41*k1[i]+a42*k2[i]+a43*k3[i])
  end
  f(@muladd(t+c3*dt),tmp,k4)
  for i in uidx
    tmp[i] = @muladd uprev[i]+dt*(a51*k1[i]+a52*k2[i]+a53*k3[i]+a54*k4[i])
  end
  f(@muladd(t+c4*dt),tmp,k5)
  for i in uidx
    tmp[i] = @muladd uprev[i]+dt*(a61*k1[i]+a62*k2[i]+a63*k3[i]+a64*k4[i]+a65*k5[i])
  end
  f(t+dt,tmp,k6)
  for i in uidx
    u[i] = @muladd uprev[i]+dt*(a71*k1[i]+a72*k2[i]+a73*k3[i]+a74*k4[i]+a75*k5[i]+a76*k6[i])
  end
  f(t+dt,u,k7)
  if integrator.opts.adaptive
    for i in uidx
      utilde[i] = @muladd uprev[i] + dt*(b1*k1[i] + b2*k2[i] + b3*k3[i] + b4*k4[i] + b5*k5[i] + b6*k6[i] + b7*k7[i])
      atmp[i] = ((utilde[i]-u[i])./@muladd(integrator.opts.abstol+max(abs(uprev[i]),abs(u[i]))*integrator.opts.reltol))
    end
    integrator.EEst = integrator.opts.internalnorm(atmp)
  end
  @pack integrator = t,dt,u,k
end

@inline function initialize!(integrator,cache::DP5ConstantCache,f=integrator.f)
  integrator.kshortsize = 4
  integrator.k = eltype(integrator.sol.k)(integrator.kshortsize)
  integrator.fsalfirst = f(integrator.t,integrator.uprev) # Pre-start fsal
end

@inline function perform_step!(integrator,cache::DP5ConstantCache,f=integrator.f)
  @unpack t,dt,uprev,u,k = integrator
  @unpack a21,a31,a32,a41,a42,a43,a51,a52,a53,a54,a61,a62,a63,a64,a65,a71,a73,a74,a75,a76,b1,b3,b4,b5,b6,b7,c1,c2,c3,c4,c5,c6 = cache
  @unpack d1,d3,d4,d5,d6,d7 = cache
  k1 = integrator.fsalfirst
  a = dt*a21
  k2 = f(@muladd(t+c1*dt),@muladd(uprev+a*k1))
  k3 = f(@muladd(t+c2*dt),@muladd(uprev+dt*(a31*k1+a32*k2)))
  k4 = f(@muladd(t+c3*dt),@muladd(uprev+dt*(a41*k1+a42*k2+a43*k3)))
  k5 = f(@muladd(t+c4*dt),@muladd(uprev+dt*(a51*k1+a52*k2+a53*k3+a54*k4)))
  k6 = f(t+dt,@muladd uprev+dt*(a61*k1+a62*k2+a63*k3+a64*k4+a65*k5))
  update = @muladd a71*k1+a73*k3+a74*k4+a75*k5+a76*k6
  u = @muladd uprev+dt*update
  integrator.fsallast = f(t+dt,u); k7 = integrator.fsallast
  if integrator.opts.adaptive
    utilde = uprev + dt*(b1*k1 + b3*k3 + b4*k4 + b5*k5 + b6*k6 + b7*k7)
    integrator.EEst = integrator.opts.internalnorm( ((utilde-u)./@muladd(integrator.opts.abstol+max(abs.(uprev),abs.(u))*integrator.opts.reltol)))
  end
  integrator.k[1] = update
  bspl = k1 - update
  integrator.k[2] = bspl
  integrator.k[3] = update - k7 - bspl
  integrator.k[4] = @muladd (d1*k1+d3*k3+d4*k4+d5*k5+d6*k6+d7*k7)
  @pack integrator = t,dt,u,k
end

@inline function initialize!(integrator,cache::DP5Cache,f=integrator.f)
  integrator.kshortsize = 4
  integrator.k = [cache.update,cache.bspl,cache.dense_tmp3,cache.dense_tmp4]
  integrator.fsalfirst = cache.k1; integrator.fsallast = cache.k7
  f(integrator.t,integrator.uprev,integrator.fsalfirst) # Pre-start fsal
end

@inline function perform_step!(integrator,cache::DP5Cache,f=integrator.f)
  @unpack t,dt,uprev,u,k = integrator
  uidx = eachindex(integrator.uprev)
  @unpack a21,a31,a32,a41,a42,a43,a51,a52,a53,a54,a61,a62,a63,a64,a65,a71,a73,a74,a75,a76,b1,b3,b4,b5,b6,b7,c1,c2,c3,c4,c5,c6 = cache.tab
  @unpack k1,k2,k3,k4,k5,k6,k7,dense_tmp3,dense_tmp4,update,bspl,utilde,tmp,atmp = cache
  @unpack d1,d3,d4,d5,d6,d7 = cache.tab
  a = dt*a21
  for i in uidx
    tmp[i] = @muladd uprev[i]+a*k1[i]
  end
  f(@muladd(t+c1*dt),tmp,k2)
  for i in uidx
    tmp[i] = @muladd uprev[i]+dt*(a31*k1[i]+a32*k2[i])
  end
  f(@muladd(t+c2*dt),tmp,k3)
  for i in uidx
    tmp[i] = @muladd uprev[i]+dt*(a41*k1[i]+a42*k2[i]+a43*k3[i])
  end
  f(@muladd(t+c3*dt),tmp,k4)
  for i in uidx
    tmp[i] = @muladd uprev[i]+dt*(a51*k1[i]+a52*k2[i]+a53*k3[i]+a54*k4[i])
  end
  f(@muladd(t+c4*dt),tmp,k5)
  for i in uidx
    tmp[i] = @muladd uprev[i]+dt*(a61*k1[i]+a62*k2[i]+a63*k3[i]+a64*k4[i]+a65*k5[i])
  end
  f(t+dt,tmp,k6)
  for i in uidx
    update[i] = @muladd a71*k1[i]+a73*k3[i]+a74*k4[i]+a75*k5[i]+a76*k6[i]
    u[i] = @muladd uprev[i]+dt*update[i]
  end
  f(t+dt,u,k7);
  if integrator.opts.adaptive
    for i in uidx
      utilde[i] = @muladd uprev[i] + dt*(b1*k1[i] + b3*k3[i] + b4*k4[i] + b5*k5[i] + b6*k6[i] + b7*k7[i])
      atmp[i] = ((utilde[i]-u[i])./@muladd(integrator.opts.abstol+max(abs(uprev[i]),abs(u[i]))*integrator.opts.reltol))
    end
    integrator.EEst = integrator.opts.internalnorm(atmp)
  end
  for i in uidx
    bspl[i] = k1[i] - update[i]
    integrator.k[3][i] = update[i] - k7[i] - bspl[i]
    integrator.k[4][i] = @muladd(d1*k1[i]+d3*k3[i]+d4*k4[i]+d5*k5[i]+d6*k6[i]+d7*k7[i])
  end
  @pack integrator = t,dt,u,k
end


@testset "uniform.jl" begin
    @testset "Constructor" begin
        @test_throws ErrorException UniformBinning(-1)
        @test_throws ErrorException UniformBinning(0)
    end

    @testset "Binning indices" begin
        # scalars
        @test_throws ArgumentError CalibrationErrors.binindex(-0.5, 10)
        @test CalibrationErrors.binindex(0, 10) == 1
        @test CalibrationErrors.binindex(0.1, 10) == 1
        @test CalibrationErrors.binindex(0.45, 10) == 5
        @test CalibrationErrors.binindex(1, 10) == 10
        @test_throws ArgumentError CalibrationErrors.binindex(1.5, 10)

        # vectors
        @test_throws ArgumentError CalibrationErrors.binindex([-0.5, 0.5], 10, Val(2))
        @test @inferred(CalibrationErrors.binindex([0, 0], 10, Val(2))) == (1, 1)
        @test @inferred(CalibrationErrors.binindex([0.1, 0], 10, Val(2))) == (1, 1)
        @test @inferred(CalibrationErrors.binindex([0.45, 0.55], 10, Val(2))) == (5, 6)
        @test @inferred(CalibrationErrors.binindex([1, 1], 10, Val(2))) == (10, 10)
        @test_throws ArgumentError CalibrationErrors.binindex([1.5, 0.5], 10, Val(2))
    end

    @testset "Basic tests" begin
        # sample predictions and targets
        nsamples = 1_000
        predictions = rand(nsamples)
        targets = rand(Bool, nsamples)

        for nbins in (1, 10, 100, 500, 1_000)
            # bin data in bins of uniform width
            bins = @inferred(
                CalibrationErrors.perform(UniformBinning(nbins), predictions, targets)
            )

            # check all bins
            for bin in bins
                # compute index of bin from average prediction
                idx = CalibrationErrors.binindex(bin.mean_predictions, nbins)

                # compute indices of all predictions in the same bin
                idxs = filter(
                    i -> idx == CalibrationErrors.binindex(predictions[i], nbins),
                    1:nsamples,
                )

                @test bin.nsamples == length(idxs)
                @test bin.mean_predictions ≈ mean(predictions[idxs])
                @test bin.proportions_targets ≈ mean(targets[idxs])
            end
        end
    end

    @testset "Basic tests ($nclasses classes)" for nclasses in (2, 10, 100)
        # sample predictions and targets
        nsamples = 1_000
        dist = Dirichlet(nclasses, 1.0)
        predictions = [rand(dist) for _ in 1:nsamples]
        targets = rand(1:nclasses, nsamples)

        for nbins in (1, 10, 100, 500, 1_000)
            # bin data in bins of uniform width
            bins = @inferred(
                CalibrationErrors.perform(UniformBinning(nbins), predictions, targets)
            )

            # check all bins
            for bin in bins
                # compute index of bin from average prediction
                idx = CalibrationErrors.binindex(bin.mean_predictions, nbins, Val(nclasses))

                # compute indices of all predictions in the same bin
                idxs = filter(
                    i ->
                        idx ==
                        CalibrationErrors.binindex(predictions[i], nbins, Val(nclasses)),
                    1:nsamples,
                )

                @test bin.nsamples == length(idxs)
                @test bin.mean_predictions ≈ mean(predictions[idxs])
                @test bin.proportions_targets ≈ proportions(targets[idxs], 1:nclasses)
            end
        end
    end

    @testset "Simple example" begin
        predictions = [[0.4, 0.1, 0.5], [0.5, 0.3, 0.2], [0.3, 0.7, 0.0]]
        targets = [1, 2, 3]

        bins = CalibrationErrors.perform(UniformBinning(2), predictions, targets)
        @test length(bins) == 2
        sort!(bins; by=x -> x.nsamples)
        @test all(bin -> sum(bin.mean_predictions) == 1, bins)
        @test all(bin -> sum(bin.proportions_targets) == 1, bins)
        for (i, idxs) in enumerate(([3], [1, 2]))
            @test bins[i].nsamples == length(idxs)
            @test bins[i].mean_predictions == mean(predictions[idxs])
            @test bins[i].proportions_targets ==
                  vec(mean(Matrix{Float64}(I, 3, 3)[:, targets[idxs]]; dims=2))
        end

        bins = CalibrationErrors.perform(UniformBinning(1), predictions, targets)
        @test length(bins) == 1
        @test bins[1].nsamples == 3
        @test bins[1].mean_predictions ≈ mean(predictions)
        @test bins[1].proportions_targets ≈ [1 / 3, 1 / 3, 1 / 3]

        predictions = [
            [0.4, 0.1, 0.5],
            [0.5, 0.3, 0.2],
            [0.3, 0.7, 0.0],
            [0.1, 0.0, 0.9],
            [0.8, 0.1, 0.1],
        ]
        targets = [1, 2, 3, 1, 2]

        bins = CalibrationErrors.perform(UniformBinning(3), predictions, targets)
        sort!(bins; by=x -> x.mean_predictions[1])
        @test length(bins) == 5
        @test all(bin -> bin.nsamples == 1, bins)
        for (i, idx) in enumerate((4, 3, 1, 2, 5))
            @test bins[i].mean_predictions == predictions[idx]
            @test bins[i].proportions_targets == Matrix{Float64}(I, 3, 3)[:, targets[idx]]
        end

        bins = CalibrationErrors.perform(UniformBinning(2), predictions, targets)
        sort!(bins; by=x -> x.mean_predictions[1])
        @test length(bins) == 4
        @test all(bin -> sum(bin.mean_predictions) == 1, bins)
        @test all(bin -> sum(bin.proportions_targets) == 1, bins)
        for (i, idxs) in enumerate(([4], [3], [1, 2], [5]))
            @test bins[i].nsamples == length(idxs)
            @test bins[i].mean_predictions == mean(predictions[idxs])
            @test bins[i].proportions_targets ==
                  vec(mean(Matrix{Float64}(I, 3, 3)[:, targets[idxs]]; dims=2))
        end

        bins = CalibrationErrors.perform(UniformBinning(1), predictions, targets)
        @test length(bins) == 1
        @test bins[1].nsamples == 5
        @test bins[1].mean_predictions ≈ mean(predictions)
        @test bins[1].proportions_targets == [0.4, 0.4, 0.2]
    end
end


using MathepiaModels
using Test
@info "Odecompartmentsmodel test"
include("Models/odecompartstest.jl")


a = rand(10^7) # Vector 1D uniforme en [0,1)

sum(a)

#Pkg.add("BenchmarkTools")

using BenchmarkTools

C_code = """
#include <stddef.h>
double c_sum(size_t n, double *X) {
    double s = 0.0;
    for (size_t i = 0; i < n; ++i) {
        s += X[i];
    }
    return s;
}
"""

const Clib = tempname()   # Haz un directorio temporario

# compila a una biblioteca compartida pipeando C_code a gcc
# (funciona sólo con gcc instalado):

open(`gcc -fPIC -O3 -msse3 -xc -shared -o $(Clib * "." * Libdl.dlext) -`, "w") do f
    print(f, C_code)
end

# define una funcion de Julia que llama a la función de C:
c_sum(X::Array{Float64}) = ccall(("c_sum", Clib), Float64, (Csize_t, Ptr{Float64}), length(X), X)

c_sum(a)

c_sum(a) ≈ sum(a) # teclea \approx y luego <TAB> para obtener el símbolo ≈

≈  # alias para la función `isapprox`

?isapprox

c_bench = @benchmark c_sum($a)

println("C: Tiempo más rápido fue $(minimum(c_bench.times) / 1e6) msec")

d = Dict()  # un diccionario a.k.a un arreglo asociativo
d["C"] = minimum(c_bench.times) / 1e6  # en milisegundos
d

using Plots
gr()

t = c_bench.times / 1e6 # tiempos en milisegundos
m, σ = minimum(t), std(t)

histogram(t, bins=500,
    xlim=(m - 0.01, m + σ),
    xlabel="milliseconds", ylabel="count", label="")

#Pkg.add("PyCall")

using PyCall

# Llama una función de bajo nivel de PyCall para obtener una lista de Python
# porque por default PyCall convertirá a un arreglo de NumPy en vez (benchmarkeamos NumPy más abajo)

apy_list = PyCall.array2py(a, 1, 1)

# obtener el "sum" interno de Python
pysum = pybuiltin("sum")

pysum(a)

pysum(a) ≈ sum(a)

py_list_bench = @benchmark $pysum($apy_list)

d["Python interno"] = minimum(py_list_bench.times) / 1e6
d

using Conda
#Conda.add("numpy")

numpy_sum = pyimport("numpy")["sum"]
apy_numpy = PyObject(a) # convierte a un arreglo de NumPy por default

py_numpy_bench = @benchmark $numpy_sum($apy_numpy)

numpy_sum(apy_list) # python thing

numpy_sum(apy_list) ≈ sum(a)

d["Python numpy"] = minimum(py_numpy_bench.times) / 1e6
d

py"""
def py_sum(a):
    s = 0.0
    for x in a:
        s = s + x
    return s
"""

sum_py = py"py_sum"

py_hand = @benchmark $sum_py($apy_list)

sum_py(apy_list)

sum_py(apy_list) ≈ sum(a)

d["Python hecho a mano"] = minimum(py_hand.times) / 1e6
d

@which sum(a)

j_bench = @benchmark sum($a)

d["Julia interno"] = minimum(j_bench.times) / 1e6
d

function mysum(A)
    s = 0.0  # s = zero(eltype(A))
    for a in A
        s += a
    end
    s
end

j_bench_hand = @benchmark mysum($a)

d["Julia hecho a mano"] = minimum(j_bench_hand.times) / 1e6
d

for (key, value) in sort(collect(d))
    println(rpad(key, 20, "."), lpad(round(value, 1), 8, "."))
end

for (key, value) in sort(collect(d), by=x->x[2])
    println(rpad(key, 20, "."), lpad(round(value, 2), 10, "."))
end

es-es/intro-to-julia-ES/plutos/12. Julia es rapido.jl


T=3
tsname="joDWT"
@testset "$tsname" begin
for t=1:T # start test loop
    m=4^t
    v1=rand(m)
    v2=rand(m,m)
    vv2=vec(v2)

    wt=wavelet(WT.haar)
    A1=joDWT(m,wt)
    A2=joDWT(m,m,wt)
    verbose && println("$tsname ($m,$m) - planned")
    @testset "$m x $m" begin
        @test isadjoint(joDWT(m))[1]
        @test islinear(joDWT(m))[1]
        @test norm(A1*v1-dwt(v1,wt))<joTol
        @test norm(A1\v1-idwt(v1,wt))<joTol
        @test norm(adjoint(A1)*v1-idwt(v1,wt))<joTol
        @test norm((adjoint(A1)*A1)*v1-v1)<joTol
        @test isadjoint(joDWT(m,m))[1]
        @test islinear(joDWT(m,m))[1]
        @test norm(A2*vv2-vec(dwt(v2,wt)))<joTol
        @test norm(A2\vv2-vec(idwt(v2,wt)))<joTol
        @test norm(adjoint(A2)*vv2-vec(idwt(v2,wt)))<joTol
        @test norm((adjoint(A2)*A2)*vv2-vv2)<joTol
    end

end # end test loop
end


module SolveDSGE

include("types.jl")
include("extra_functions.jl")
include("givens.jl")
include("householder.jl")
include("vec_to_vech.jl")
include("vech_to_vec.jl")
include("derivative.jl")
include("hessian.jl")
include("reorder_complex_schur.jl")
include("check_model.jl")
include("solve_re.jl")

export derivative,
       hessian,
       solve_re,
       check_model_form,
       convert_second_order

export Blanchard_Kahn_Form,
       Blanchard_Kahn_Soln,
       Klein_Form,
       Klein_Soln,
       Structural_Form,
       Structural_Soln,
       Sims_Form,
       Sims_Soln,
       Gomme_Klein_Form,
       Gomme_Klein_Soln,
       Lombardo_Sutherland_Form,
       Lombardo_Sutherland_Soln

end


### A Pluto.jl notebook ###
# v0.16.1

using Markdown
using InteractiveUtils

# ╔═╡ 09bee220-1249-11ec-3e2a-4fdde9976a59
begin
	import Pkg
    Pkg.activate(mktempdir())
	Pkg.add(url="https://github.com/brendanjohnharris/Catch22.jl#main")
	using Catch22
	Pkg.add(url="https://github.com/tk3369/YeoJohnsonTrans.jl")
	Pkg.develop(path="C:\\Users\\Brendan\\OneDrive - The University of Sydney (Students)\\Honours\\Code\\ParameterInference.jl")
	Pkg.add(path="C:\\Users\\Brendan\\OneDrive - The University of Sydney (Students)\\Honours\\Code\\ParameterInference.jl")
	using ParameterInference, Plots, StatsBase, Statistics, LinearAlgebra, ParameterInference.NonstationaryProcesses, MultivariateStats
	gr(); fourseas!(); # Plots initialisation
	nothing
end

# ╔═╡ 1cbf17f7-8322-4e3c-a895-163a48332205
md"# Baseline PCA"

# ╔═╡ 2f0cb4e5-4424-42d9-8eb7-e8bf128575ad
md"""
We will construct a dataset that demonstrates the impact of dependencies between features on PCA estimates.
"""

# ╔═╡ 0f10dc15-9153-4d8c-a5e5-c0318c459fa1
md"## Parameters"

# ╔═╡ 230c66e5-7898-48af-9544-439677bee137
md"For a simple example, choose to have two 'parameters' that drive our manually constructed 'features': 𝛼 and 𝛽. Both will vary as random Gaussians.:"

# ╔═╡ fee5a67f-1678-45fc-83b3-3ffc187c19b3
nₛ = 1000;

# ╔═╡ c3e707a4-1e35-4fe0-8e6d-88d0eca27fde
𝜂 = 0.25;

# ╔═╡ 5146ae95-d551-49ec-888a-71fc7bf65d14
𝛼, 𝛽 = randn(nₛ, 2) |> eachcol;

# ╔═╡ b646056a-7c34-4cb3-aac9-d255646f6b75
plotcov(x) = covarianceimage(1:5, x, colormode=:raw, docluster=false, colorbar_title="", title="\$|\\Sigma^2|\$");

# ╔═╡ 2fbff3c4-9377-4624-90d6-d4ea7c8da64e
md"# Feature scenario #1"

# ╔═╡ 404bb5c9-74ab-43a9-8791-67ac23b0c0f8
md"""Next, we construct 5 features according to the parameters.
In our first scenario, these are:

| Feature      	| Description  	|
| ----------- 	| ----------- 	|
| 𝑓₁ -- 𝑓₄      | Sensitive to 𝛼, but intercorrelated  |
| 𝑓₅   		 | Sensitive to 𝛽  |
"""

# ╔═╡ 4f01861c-dd52-462f-b340-ceecd3385c8e
md"""### High-dimensional baseline
In which both 𝛼 and 𝛽 are varying:
"""

# ╔═╡ 92670bd1-ec33-4590-9a53-ceca6f25ac0f
𝑓ₕ₁ = hcat(𝛼, 𝛼, 𝛼, 𝛼, 𝛽)';

# ╔═╡ 8629c505-a1e4-4c73-b2d4-44ecc858373d
plotcov(𝑓ₕ₁)

# ╔═╡ c3c3d6c3-49a0-40a6-b41b-dbdfa20a7d6b
md"""### Test data
In which 𝛼 is constrained:
"""

# ╔═╡ 9247ba48-f242-4479-a92d-e994b1610279
𝑓₁ = hcat(𝜂.*[𝛼, 𝛼, 𝛼, 𝛼]..., 𝛽)';

# ╔═╡ c973721b-c20b-4573-a8fa-a9a9282fc87d
plotcov(𝑓₁)

# ╔═╡ 4032a036-842b-462b-b75d-ebaaaf4aded9
md"### True parameters"

# ╔═╡ 6a957b7d-8fff-46c5-b441-01f8b9d45eb7
scatter(𝛽, 𝜂.*𝛼, aspect_ratio=:equal, ylabel="𝛼", xlabel="𝛽", title="True parameter variation",  markersize=2, markercolor=:gray)

# ╔═╡ 5686638b-ba3f-4691-9291-d91dfa510aa7
md"### PCA"

# ╔═╡ 76e702f7-521c-4445-aa6c-de28215d8d44
𝛴² = cov(𝑓₁');

# ╔═╡ b17a0c26-ce25-4e11-84dd-0d4b86371ae5
𝑃 = eigvecs(𝛴²);

# ╔═╡ 69e0799e-18f4-4f19-9bae-44202c2bba5d
𝑓′₁ = 𝑃'*𝑓₁;

# ╔═╡ cf454ccf-e505-4df9-a3cf-4afba2211d4a
scatter(𝑓′₁[5, :], 𝑓′₁[4, :], aspect_ratio=:equal, ylabel="𝛼′", xlabel="𝛽′", title="PCA estimate"); # Scatter significant PCs

# ╔═╡ 8013d3a2-a011-456b-9bfe-3ec781c31aba
scatter!(𝛽, 𝜂.*𝛼, aspect_ratio=:equal, markersize=2, markercolor=:gray)

# ╔═╡ ce024bfd-b892-4a31-a0d1-1d4befc7b80e
md"### Standardised PCA"

# ╔═╡ 542aec21-e623-4e53-ae4b-a7f240bfcb08
𝑓̂₁ = 𝑓₁./std(𝑓₁, dims=2);

# ╔═╡ 1daa139a-23c0-4995-bf73-abfbb7c518dd
𝛴̂² = cov(𝑓̂₁');

# ╔═╡ 705068d7-033f-49d2-bb59-4b40aa514642
𝑃̂ =  fit(MultivariateStats.PCA, 𝑓̂₁).proj; #eigvecs(𝛴̂²);

# ╔═╡ f1a685e5-ab99-4a61-a2d9-41ebf33372b0
𝑓̂′₁ = 𝑃̂'*𝑓₁;

# ╔═╡ 3a7af0db-8e3d-4100-a7c2-dcd37253a9ce
scatter(𝑓̂′₁[2, :], 𝑓̂′₁[1, :], aspect_ratio=:equal, ylabel="𝛼′", xlabel="𝛽′", title="Standardised PCA estimate");

# ╔═╡ 7a28c225-bab8-40ee-906e-0c5349accc84
scatter!(𝛽, 𝜂.*𝛼, aspect_ratio=:equal, markersize=2, markercolor=:gray)

# ╔═╡ 4d619965-6eb7-4528-804c-00bbf1572e60
md"""### Baseline PCA
In this case, whitening:"""

# ╔═╡ 7ddd5530-7590-42a7-8e66-54432f717eec
𝛴²ₕ = cov(𝑓ₕ₁');

# ╔═╡ a95a3d4c-cb8b-4840-bc67-57fb2be00065
m1 = fit(MultivariateStats.PCA, 𝑓ₕ₁; maxoutdim=5, pratio=1.0)

# ╔═╡ 688dc093-c173-4eb4-ad12-4ff2e899b9f9
𝑃ₕ = m1.proj; # eigvecs(𝛴²ₕ)

# ╔═╡ f4081452-9314-47a4-948d-2813bda53507
𝛬ₕ = sqrt(diagm(m1.prinvars)); # Diagonal(sqrt.(abs.(eigvals(𝛴²ₕ))));

# ╔═╡ 5b853331-88c7-4b7f-b26e-8e0f963a5ed5
𝑓′′₁ = (𝑃ₕ*inv(𝛬ₕ))'*𝑓₁;

# ╔═╡ 1ee56732-671e-447e-adb8-cce45902b935
scatter(𝑓′′₁[2, :], .-𝑓′′₁[1, :], aspect_ratio=:equal, ylabel="𝛼′", xlabel="𝛽′");

# ╔═╡ 70d2fe5f-7f7d-4891-8490-5c508c21c6f6
scatter!(𝛽, 𝜂.*𝛼, aspect_ratio=:equal, markersize=2, markercolor=:gray)

# ╔═╡ bb52a2b3-e856-4785-abb9-78ba15382d11
md"# Feature scenario #2"

# ╔═╡ 70d2f196-1604-4cf4-8e38-05dd8c099c83
md"""What if features one to four are sensitive to both 𝛼 and 𝛽? Then the effect of constraining 𝛼 manifests as an increase in correlation between features. For our second scenario, the features are then:

| Feature      	| Description  	|
| ----------- 	| ----------- 	|
| 𝑓₁ -- 𝑓₄      | Sensitive to 𝛼 and 𝛽, but intercorrelated  |
| 𝑓₅   		 | Sensitive to 𝛽  |
"""

# ╔═╡ 009c662a-9940-43ef-b0c9-d55560e6f49b
md"### High-dimensional baseline"

# ╔═╡ 46e7d2a0-8b62-4112-92df-74538f5f16d8
𝑓ₕ₂ = hcat(𝛼+𝛽/4, 𝛼+𝛽/4, 𝛼+𝛽/4, 𝛼+𝛽/4, 𝛽)';

# ╔═╡ 595e5022-9099-45ab-b5b0-b92fedb96ba8
plotcov(𝑓ₕ₂./std(𝑓ₕ₂, dims=2))

# ╔═╡ 0d062a65-adca-42c4-8b1c-84095d5335c5
md"### Test data"

# ╔═╡ 0f6b9065-55ad-4b1e-bcc3-2e177a0e331d
𝑓₂ = hcat((𝜂.*[𝛼, 𝛼, 𝛼, 𝛼] .+ [𝛽, 𝛽, 𝛽, 𝛽]./4)..., 𝛽)';

# ╔═╡ fd506ba8-9a5b-4235-bc43-b730522dce4f
plotcov(𝑓₂./std(𝑓₂, dims=2))

# ╔═╡ 58083319-d649-4c2d-b4b2-c2df7c49ee46
md"### True parameters"

# ╔═╡ 2a9a4929-8755-4237-a454-fda125d548ca
scatter(𝛽, 𝜂.*𝛼 .+ 𝛽./4, aspect_ratio=:equal, ylabel="𝛼", xlabel="𝛽", title="True parameter variation",  markersize=2, markercolor=:gray)

# ╔═╡ 1762e05f-7eec-4411-865a-eba47531a49c
md"### Standardised PCA"

# ╔═╡ feeaa835-28e0-4a24-905b-53347560ec9a
𝑓̂₂ = 𝑓₂./std(𝑓₂, dims=2);

# ╔═╡ 2358f79f-f44d-4e93-b4da-75e7be715882
𝛴̂²₂ = cov(𝑓̂₂');

# ╔═╡ d9fca85a-63b6-4ab5-ab0e-6cd5e68c4e06
𝑃̂₂ = fit(MultivariateStats.PCA, 𝑓̂₂).proj; # eigvecs(𝛴̂²₂);

# ╔═╡ d700b4ae-8bd5-4437-b515-c21d3da70543
𝑓̂′₂ = 𝑃̂₂'*𝑓₂;

# ╔═╡ 06e226a6-0f0f-4ef6-b916-04d51a3789c8
scatter(𝑓̂′₂[1, :], .-𝑓̂′₂[2, :], aspect_ratio=:equal, ylabel="𝛼′", xlabel="𝛽′", title="Standardised PCA estimate");

# ╔═╡ ea513127-e9bb-40a7-9ec0-2beee8fa6179
scatter!(𝛽, 𝜂.*𝛼 .+ 𝛽./4, markersize=2, markercolor=:gray)

# ╔═╡ f0a7f529-19ac-44f7-8ab1-8d72d596a0ef
md"""### Baseline PCA"""

# ╔═╡ f946e3d7-15b9-45a9-951b-b3d77fc589f8
𝛴²ₕ₂ = cov(𝑓ₕ₂');

# ╔═╡ 652b998d-a6b0-4bb7-908e-c3dc7a82764c
m2 = fit(MultivariateStats.PCA, 𝑓ₕ₂);

# ╔═╡ d192149c-1ef4-4761-9e78-879a97abee57
𝑃ₕ₂ = m2.proj; #eigvecs(𝛴²ₕ₂);

# ╔═╡ 14d5e08c-0808-4352-835e-c173c53fd8d7
𝛬ₕ₂ = sqrt(diagm(m2.prinvars)); # Diagonal(sqrt.(abs.(eigvals(𝛴²ₕ₂))))

# ╔═╡ 772d260a-72b3-4c9e-8e65-6368508cfb52
𝑓′′₂ = (𝑃ₕ₂*inv(𝛬ₕ₂))'*𝑓₂;

# ╔═╡ 0f83a03c-890d-4771-b83f-16d8d5db81e9
cov(𝑓′′₂,dims=2)

# ╔═╡ b9a8e7e4-5f9b-462d-845e-b873f50bfc90
scatter(𝑓′′₂[2, :], .-𝑓′′₂[1, :], aspect_ratio=:equal, ylabel="𝛼′", xlabel="𝛽′");

# ╔═╡ f1df27ea-0537-4f4b-afe7-c0692ca7b26d
scatter!(𝛽, 𝜂.*𝛼 .+ 𝛽./4, markersize=2, markercolor=:gray)

# ╔═╡ 2f987bde-a87a-43a8-8653-0c5812c56a70
md"# Feature scenario #3"

# ╔═╡ 1d18e8c7-7eff-4b75-ad67-e61bc22e12ea
md"""What if each feature has some intrinsic noise? This requires incorporating low-dimensional variance into the estimate.
To show this:

| Feature      	| Description  	|
| ----------- 	| ----------- 	|
| 𝑓₁ -- 𝑓₃      | Sensitive to 𝛼, but intercorrelated and with small noise |
| 𝑓₄   		 | Sensitive to 𝛽, but with small noise  |
| 𝑓₅   		 | Small noise only  |
"""

# ╔═╡ 3821aa68-e3b3-4764-a7ff-0b6b0e6bdc3e
𝜂₃ = 0.25;

# ╔═╡ 9d0c3f2b-8cf4-41f9-8fe8-5cd9a19dbd88
𝑓ₕ₃ = hcat(repeat([𝛼.+𝜂₃.*randn(nₛ)],3)..., 𝛽.+𝜂₃.*randn(nₛ), (𝜂₃.-0.01).*randn(nₛ))';

# ╔═╡ b0d86eb8-6f94-424a-adb2-eae8e40be1fd
plotcov(𝑓ₕ₃)

# ╔═╡ e507beed-984e-4f8b-a59b-bb3dcff1416a
md"""### Zero-dimensional baseline
In which no parameters are varying:
"""

# ╔═╡ ea23a142-1c36-4430-a598-2fcf0f2e3b66
𝑓₀ = hcat(repeat([𝜂₃.*randn(nₛ)], 3)..., 𝜂₃.*randn(nₛ), 𝜂₃.*randn(nₛ))';

# ╔═╡ fe0902aa-b55e-4b72-a35e-77db673747c9
plotcov(𝑓₀)

# ╔═╡ 9b92548e-ae54-48de-a5c9-4be88ae11228
md"### Test data"

# ╔═╡ 33167762-e2c6-4aba-98b7-58e28ee35ccb
𝑓₃ = hcat(repeat([0.25.*𝛼 .+ 𝜂₃.*randn(nₛ)], 3)..., 𝛽.+𝜂₃.*randn(nₛ), 𝜂₃.*randn(nₛ))';

# ╔═╡ 35dcd411-a342-4ee2-9146-f45ad39d6341
plotcov(𝑓₃)

# ╔═╡ d2404839-53bf-49d3-a9bf-a0ea81471c43
md"### True parameters"

# ╔═╡ e64a04de-9257-4bbb-84c4-007c532d9c72
scatter(𝛽, 0.25.*𝛼, aspect_ratio=:equal, ylabel="𝛼", xlabel="𝛽", title="True parameter variation",  markersize=2, markercolor=:gray)

# ╔═╡ 710d51ea-8196-4a7f-91fe-1a9bd3f06d01
md"""
### Standardised PCA
"""

# ╔═╡ a4c8677f-6bdc-4bfe-adcc-3130f45b7aaa
𝑓̂₃ = 𝑓₃./std(𝑓₃, dims=2);

# ╔═╡ 5ad01546-7daa-441f-bf2c-60e2a6e6f0c7
𝛴̂²₃ = cov(𝑓̂₃');

# ╔═╡ 36039419-6df3-45f2-ae2c-77cf6e3b405d
𝑃̂₃ = eigvecs(𝛴̂²₃);

# ╔═╡ d046eb57-d4cb-4eda-9a81-00a234af1d57
𝑓̂′₃ = 𝑃̂₃'*𝑓₃;

# ╔═╡ 4e9b0004-da7a-4cb3-afe9-c3365d37471a
scatter(𝑓̂′₃[4, :], 𝑓̂′₃[5, :], aspect_ratio=:equal, ylabel="𝛼′", xlabel="𝛽′", title="Standardised PCA estimate");

# ╔═╡ 5e2281e9-360b-41dd-9f4f-38e14d8aea86
scatter!(𝛽, 0.25.*𝛼, markersize=2, markercolor=:gray)

# ╔═╡ fa43e761-0b22-4903-aa6c-47050d8c27c0
scatter(𝛼, 𝑓̂′₃[5, :])

# ╔═╡ bcd85e11-8bf0-4968-a528-8b56ce14db5b
md"""### Baseline PCA"""

# ╔═╡ 3fe79e11-947e-46e6-a5ef-1bf1b537c857
md"""
The full baseline-corrected covariance is:

$$\Sigma^{\prime 2} = (P \Lambda_h^{-1})^T \left(\Sigma^2 - \Sigma_0^2\right)(P \Lambda_h^{-1}).$$

But this is not a linear transformation of feature space, so some approximations must (?) be made.

An easy one is to treat the zero-dimensional covariance as diagonal. In that case:

$$\Sigma^2 - \Sigma_0^2 \approx V^T \Sigma^2 V,$$
where:

$$V = \sqrt{I - D(\Sigma^2_0) D(\Sigma^{-2})},$$
And $D(A)$ is the diagonal of $A$. The resulting approximation is simply a linear rescaling. Importantly, it satisfies the condition that $V \to 0$ as $\Sigma^2 \to \Sigma^2_0$.

There are other approximations that also allow covariances to vanish in the same limit...
"""

# ╔═╡ 15970d81-af01-443e-b54a-bda505b1b3be
md"""
For now, we will just look at a heuristic approximation:

$$V = I - D(\Sigma') D(\Sigma_0')$$

which gives a total transformation of:

$$F^\prime = FP\frac{I - D(\Sigma^\prime)^{-1}D(\Sigma^\prime_0)}{\Lambda_h - D(\Sigma^\prime_0)}$$

where $P$ is the projection matrix of the high-dimensional baseline and the primes indicate application of $P$.

"""

# ╔═╡ f771887a-77f0-4a4d-9e4c-7e988201197f
𝛴²ₕ₃ = cov(𝑓ₕ₃');

# ╔═╡ b636df98-f935-46b4-bb08-be96cee18884
𝛴²₀ = cov(𝑓₀');

# ╔═╡ c43d6053-4eeb-450a-bca4-bd5d58b179c1
𝛴²₃ = cov(𝑓₃');

# ╔═╡ beafaf38-3753-4e9f-bf6d-69b98439fe18
m3 = fit(MultivariateStats.PCA, 𝑓ₕ₃)

# ╔═╡ fd00b9e4-f2fc-4571-8596-b1e53f8d7286
𝑃ₕ₃ = m3.proj;#eigvecs(𝛴²ₕ₃);

# ╔═╡ 9d13988d-e625-42b5-8e89-86d519027912
𝛬ₕ₃ = sqrt(diagm(m3.prinvars)); #Diagonal(sqrt.(abs.(eigvals(𝛴²ₕ₃))))

# ╔═╡ 4f5c167d-dbde-4fed-bd2d-63b0100e2d1a
𝛴²ₕ′ = diagm(m3.prinvars)

# ╔═╡ 82841686-b078-4ba2-b5ce-d748070437d2
𝛴²₀′ = cov(𝑃ₕ₃'*𝑓₀, dims=2)

# ╔═╡ 8dded5c7-f4ea-4ee5-b0b9-59b21c7f3097
𝛴²′ = cov(𝑃ₕ₃'*𝑓₃, dims=2)

# ╔═╡ cb77a71b-7cba-489d-a51b-690cdd281a4c
𝑉 = I - inv(sqrt(Diagonal(𝛴²′)))*sqrt(Diagonal(𝛴²₀′));

# ╔═╡ b00f965a-18a1-4ccf-85a9-ae79d588cab6
𝑆 = 𝑉/(𝛬ₕ₃ - sqrt(Diagonal(𝛴²₀′)))

# ╔═╡ 0ae23619-0ce3-4e32-9a11-f9a444235c70
𝑓′′₃ = 𝑆'*𝑃ₕ₃'*𝑓₃;

# ╔═╡ ddb9528f-b57c-4333-9347-8a56aa6890c4
#f′′₃ = (𝑓′′₃./std(𝑓′′₃, dims=2))*((sqrt(Diagonal(𝛴²′)) - sqrt(Diagonal(𝛴²₀′)))/())

# ╔═╡ c3067e6d-dfaf-4938-8bbf-4b51514fd286
𝑓′′₀ = (𝑃ₕ₃*inv(𝛬ₕ₃)*𝑉)'*𝛴²₀*(𝑃ₕ₃*inv(𝛬ₕ₃)*𝑉); # Not 0 because the transform is a strong approximation

# ╔═╡ a2f4c245-031b-4f1c-9f44-b3947cbc9795
cov(𝑓′′₃, dims=2)

# ╔═╡ ad7aac3b-303a-4ca6-a44f-f11cba8d115f
𝑓′′₀

# ╔═╡ dbdcd6c4-102d-41cc-b218-100f74355b85
scatter(𝑓′′₃[2, :], .-𝑓′′₃[1, :], aspect_ratio=:equal, ylabel="𝛼′", xlabel="𝛽′", title="Baseline PCA estimate", left_margin=5Plots.mm);

# ╔═╡ 9d6334c1-7107-4fbb-b383-9a37fc56de74
scatter!(.-𝛽, .-0.25.*𝛼, markersize=1, markercolor=:gray)

# ╔═╡ 07e9a417-5adb-4182-a98d-fd618feaf675
scatter(𝛽, 𝑓′′₃[2, :]) # Compairwise to true parameters

# ╔═╡ 008512c1-3a44-42b4-8695-2a9f4d88ef01
scatter(𝛼, 𝑓′′₃[1, :]) # Compairwise to true parameters

# ╔═╡ 2826e49a-5566-4aa7-be9a-73799dfbebc7
scatter(𝑓̂′₃[4, :], 𝑓′′₃[4, :]); # Compare to standardised PCA estimate. Similar, which is good, but properly scaled relative to other PCs (as above).

# ╔═╡ c79d0914-d0db-4f5c-b1f2-a6154d063e78

# ╔═╡ eaae76e1-9326-432a-8fc5-b03c3d93e75a

# ╔═╡ b7a9e831-6ab4-4afd-a5da-9d3b5414245b
fₕ₃ = 𝑓ₕ₃ .+eps().*randn(size(𝑓ₕ₃))

# ╔═╡ 70109d54-cfcd-46c5-8c67-a3b96286ee09
f₃ = 𝑓₃ .+eps().*randn(size(𝑓₃))

# ╔═╡ 9970c6e2-a2c9-442e-8fca-17b8a5c360b9
𝑃ₕ₄ = eigvecs(𝛴²ₕ₃)

# ╔═╡ 7408597f-8bc6-493a-80f6-79a5201d625b
#Λ = Diagonal(sqrt.(abs.(eigvals(𝛴²ₕ₃))))

# ╔═╡ a031cdd3-7e92-45d8-8037-e6695cd588c2
MultivariateStats.PCA

# ╔═╡ 816e841b-3f3b-4b8c-8c4b-4cd93c685ab7
m = fit(MultivariateStats.PCA, 𝑓ₕ₃; maxoutdim=5, pratio=1.0)

# ╔═╡ 923b9503-5a8e-4dab-984d-6ca07d625bb4
Σₕ = cov(fₕ₃, dims=2)

# ╔═╡ 9938d1b8-8f80-45b0-bdff-48cd74e78fa4
P = m.proj

# ╔═╡ 4771733f-83a5-4c13-b0ea-cb54241f7e19
Λ² = diagm(eigvals(Σₕ))

# ╔═╡ d8847644-4d5b-4113-8870-e8c3dba7e1a2
Λ²[Λ².<0] .= eps()

# ╔═╡ 5463ac9d-9bf7-4363-abb9-d1de01a88681
Λ = sqrt(diagm(m.prinvars))

# ╔═╡ 39871429-722a-4745-9a67-62c2961fb268
ffh = P'*𝑓ₕ₃

# ╔═╡ 31a016af-82fa-41a1-b2e8-8adea326c11c
ff = P'*𝑓₃

# ╔═╡ 31cfab8a-e219-484c-9428-1fef8ea16ca5
ff0 = P'*𝑓₀

# ╔═╡ c64a1c15-824a-49d2-ac0f-10573a1472c5
l0 = cov(ff0, dims=2)

# ╔═╡ 8c9e6bd3-c548-4126-a6fb-a100e91f8747
l = cov(ff, dims=2)

# ╔═╡ 24425756-1e29-4e78-9398-6ef1820940b8
lh = Λ^2

# ╔═╡ 4a74c5b1-d19b-4add-a0ca-3a239d8a7c75
l*(I - Diagonal(l0)/Diagonal(l))/(Λ - l0)

# ╔═╡ 3914f4ec-4b70-4f05-9ec1-a9f70223bfd8
inv(Λ)[3, 3]*l[3, 3]*(1-l0[3, 3]/l[3, 3])

# ╔═╡ 94cf8ec3-5c81-4874-a331-aa0f15563808
a = (P*inv(Λ))'*(𝛴²ₕ₃)*(P*inv(Λ))

# ╔═╡ ffb3c3bb-827e-4173-90af-231bc45fef11
𝛴²₃

# ╔═╡ 2a1d988a-21fd-43b1-9701-dcfcf8f3d09b
A = (P*sqrt((I - Diagonal(l0)/Diagonal(l))/(Λ - l0)))'*(𝛴²₃)*P*sqrt((I - Diagonal(l0)/Diagonal(l))/(Λ - l0))

# ╔═╡ f51dc362-e8ab-494b-94f9-7ec63262afe4
B = (P*inv(Λ))'*𝛴²₀*P*inv(Λ)

# ╔═╡ 9ae7ea6f-6f70-4d92-9009-4ed3ccde3945
C = Real.(((I - sqrt(Diagonal(l0))/sqrt(Diagonal(l)))/(sqrt(Λ) - sqrt(l0))))'*ff

# ╔═╡ 5a0bd5e0-e3fb-47dc-8be2-fec23ea865a9
cov(C, dims=2)

# ╔═╡ 2df2e495-5170-4db3-904b-0c272f69b182
D = Real.(sqrt((I - Diagonal(l0)/Diagonal(l))/(Λ - l0)))'*ff0

# ╔═╡ 53142c10-19e9-42ba-b06d-85e326c8c575
cov(D, dims=2)

# ╔═╡ 008b10fc-8e80-4783-ae37-1f18147d6f55
scatter(C[2, :], .-C[3, :], aspect_ratio=:equal, ylabel="𝛼′", xlabel="𝛽′", title="Standardised PCA estimate", left_margin=5Plots.mm)

# ╔═╡ c43d7521-0e3d-403c-83b6-d9cce7675931
scatter!(.-𝛽, .-0.1.*𝛼, markersize=1, markercolor=:gray)

# ╔═╡ Cell order:
# ╠═09bee220-1249-11ec-3e2a-4fdde9976a59
# ╟─1cbf17f7-8322-4e3c-a895-163a48332205
# ╟─2f0cb4e5-4424-42d9-8eb7-e8bf128575ad
# ╟─0f10dc15-9153-4d8c-a5e5-c0318c459fa1
# ╟─230c66e5-7898-48af-9544-439677bee137
# ╠═fee5a67f-1678-45fc-83b3-3ffc187c19b3
# ╠═c3e707a4-1e35-4fe0-8e6d-88d0eca27fde
# ╠═5146ae95-d551-49ec-888a-71fc7bf65d14
# ╠═b646056a-7c34-4cb3-aac9-d255646f6b75
# ╟─2fbff3c4-9377-4624-90d6-d4ea7c8da64e
# ╟─404bb5c9-74ab-43a9-8791-67ac23b0c0f8
# ╟─4f01861c-dd52-462f-b340-ceecd3385c8e
# ╠═92670bd1-ec33-4590-9a53-ceca6f25ac0f
# ╠═8629c505-a1e4-4c73-b2d4-44ecc858373d
# ╟─c3c3d6c3-49a0-40a6-b41b-dbdfa20a7d6b
# ╠═9247ba48-f242-4479-a92d-e994b1610279
# ╠═c973721b-c20b-4573-a8fa-a9a9282fc87d
# ╟─4032a036-842b-462b-b75d-ebaaaf4aded9
# ╠═6a957b7d-8fff-46c5-b441-01f8b9d45eb7
# ╟─5686638b-ba3f-4691-9291-d91dfa510aa7
# ╠═76e702f7-521c-4445-aa6c-de28215d8d44
# ╠═b17a0c26-ce25-4e11-84dd-0d4b86371ae5
# ╠═69e0799e-18f4-4f19-9bae-44202c2bba5d
# ╠═cf454ccf-e505-4df9-a3cf-4afba2211d4a
# ╠═8013d3a2-a011-456b-9bfe-3ec781c31aba
# ╟─ce024bfd-b892-4a31-a0d1-1d4befc7b80e
# ╠═542aec21-e623-4e53-ae4b-a7f240bfcb08
# ╠═1daa139a-23c0-4995-bf73-abfbb7c518dd
# ╠═705068d7-033f-49d2-bb59-4b40aa514642
# ╠═f1a685e5-ab99-4a61-a2d9-41ebf33372b0
# ╠═3a7af0db-8e3d-4100-a7c2-dcd37253a9ce
# ╠═7a28c225-bab8-40ee-906e-0c5349accc84
# ╟─4d619965-6eb7-4528-804c-00bbf1572e60
# ╠═7ddd5530-7590-42a7-8e66-54432f717eec
# ╠═a95a3d4c-cb8b-4840-bc67-57fb2be00065
# ╠═688dc093-c173-4eb4-ad12-4ff2e899b9f9
# ╠═f4081452-9314-47a4-948d-2813bda53507
# ╠═5b853331-88c7-4b7f-b26e-8e0f963a5ed5
# ╠═1ee56732-671e-447e-adb8-cce45902b935
# ╠═70d2fe5f-7f7d-4891-8490-5c508c21c6f6
# ╟─bb52a2b3-e856-4785-abb9-78ba15382d11
# ╟─70d2f196-1604-4cf4-8e38-05dd8c099c83
# ╟─009c662a-9940-43ef-b0c9-d55560e6f49b
# ╠═46e7d2a0-8b62-4112-92df-74538f5f16d8
# ╠═595e5022-9099-45ab-b5b0-b92fedb96ba8
# ╟─0d062a65-adca-42c4-8b1c-84095d5335c5
# ╠═0f6b9065-55ad-4b1e-bcc3-2e177a0e331d
# ╠═fd506ba8-9a5b-4235-bc43-b730522dce4f
# ╟─58083319-d649-4c2d-b4b2-c2df7c49ee46
# ╠═2a9a4929-8755-4237-a454-fda125d548ca
# ╟─1762e05f-7eec-4411-865a-eba47531a49c
# ╠═feeaa835-28e0-4a24-905b-53347560ec9a
# ╠═2358f79f-f44d-4e93-b4da-75e7be715882
# ╠═d9fca85a-63b6-4ab5-ab0e-6cd5e68c4e06
# ╠═d700b4ae-8bd5-4437-b515-c21d3da70543
# ╠═06e226a6-0f0f-4ef6-b916-04d51a3789c8
# ╠═ea513127-e9bb-40a7-9ec0-2beee8fa6179
# ╟─f0a7f529-19ac-44f7-8ab1-8d72d596a0ef
# ╠═f946e3d7-15b9-45a9-951b-b3d77fc589f8
# ╠═652b998d-a6b0-4bb7-908e-c3dc7a82764c
# ╠═d192149c-1ef4-4761-9e78-879a97abee57
# ╠═14d5e08c-0808-4352-835e-c173c53fd8d7
# ╠═772d260a-72b3-4c9e-8e65-6368508cfb52
# ╠═0f83a03c-890d-4771-b83f-16d8d5db81e9
# ╠═b9a8e7e4-5f9b-462d-845e-b873f50bfc90
# ╠═f1df27ea-0537-4f4b-afe7-c0692ca7b26d
# ╟─2f987bde-a87a-43a8-8653-0c5812c56a70
# ╟─1d18e8c7-7eff-4b75-ad67-e61bc22e12ea
# ╠═3821aa68-e3b3-4764-a7ff-0b6b0e6bdc3e
# ╠═9d0c3f2b-8cf4-41f9-8fe8-5cd9a19dbd88
# ╠═b0d86eb8-6f94-424a-adb2-eae8e40be1fd
# ╟─e507beed-984e-4f8b-a59b-bb3dcff1416a
# ╠═ea23a142-1c36-4430-a598-2fcf0f2e3b66
# ╠═fe0902aa-b55e-4b72-a35e-77db673747c9
# ╟─9b92548e-ae54-48de-a5c9-4be88ae11228
# ╠═33167762-e2c6-4aba-98b7-58e28ee35ccb
# ╠═35dcd411-a342-4ee2-9146-f45ad39d6341
# ╟─d2404839-53bf-49d3-a9bf-a0ea81471c43
# ╠═e64a04de-9257-4bbb-84c4-007c532d9c72
# ╟─710d51ea-8196-4a7f-91fe-1a9bd3f06d01
# ╠═a4c8677f-6bdc-4bfe-adcc-3130f45b7aaa
# ╠═5ad01546-7daa-441f-bf2c-60e2a6e6f0c7
# ╠═36039419-6df3-45f2-ae2c-77cf6e3b405d
# ╠═d046eb57-d4cb-4eda-9a81-00a234af1d57
# ╠═4e9b0004-da7a-4cb3-afe9-c3365d37471a
# ╠═5e2281e9-360b-41dd-9f4f-38e14d8aea86
# ╠═fa43e761-0b22-4903-aa6c-47050d8c27c0
# ╟─bcd85e11-8bf0-4968-a528-8b56ce14db5b
# ╟─3fe79e11-947e-46e6-a5ef-1bf1b537c857
# ╠═15970d81-af01-443e-b54a-bda505b1b3be
# ╠═f771887a-77f0-4a4d-9e4c-7e988201197f
# ╠═b636df98-f935-46b4-bb08-be96cee18884
# ╠═c43d6053-4eeb-450a-bca4-bd5d58b179c1
# ╠═beafaf38-3753-4e9f-bf6d-69b98439fe18
# ╠═fd00b9e4-f2fc-4571-8596-b1e53f8d7286
# ╠═9d13988d-e625-42b5-8e89-86d519027912
# ╠═4f5c167d-dbde-4fed-bd2d-63b0100e2d1a
# ╠═82841686-b078-4ba2-b5ce-d748070437d2
# ╠═8dded5c7-f4ea-4ee5-b0b9-59b21c7f3097
# ╠═cb77a71b-7cba-489d-a51b-690cdd281a4c
# ╠═b00f965a-18a1-4ccf-85a9-ae79d588cab6
# ╠═0ae23619-0ce3-4e32-9a11-f9a444235c70
# ╠═ddb9528f-b57c-4333-9347-8a56aa6890c4
# ╠═c3067e6d-dfaf-4938-8bbf-4b51514fd286
# ╠═a2f4c245-031b-4f1c-9f44-b3947cbc9795
# ╠═ad7aac3b-303a-4ca6-a44f-f11cba8d115f
# ╠═dbdcd6c4-102d-41cc-b218-100f74355b85
# ╠═9d6334c1-7107-4fbb-b383-9a37fc56de74
# ╠═07e9a417-5adb-4182-a98d-fd618feaf675
# ╠═008512c1-3a44-42b4-8695-2a9f4d88ef01
# ╠═2826e49a-5566-4aa7-be9a-73799dfbebc7
# ╠═c79d0914-d0db-4f5c-b1f2-a6154d063e78
# ╠═eaae76e1-9326-432a-8fc5-b03c3d93e75a
# ╠═b7a9e831-6ab4-4afd-a5da-9d3b5414245b
# ╠═70109d54-cfcd-46c5-8c67-a3b96286ee09
# ╠═9970c6e2-a2c9-442e-8fca-17b8a5c360b9
# ╠═7408597f-8bc6-493a-80f6-79a5201d625b
# ╠═a031cdd3-7e92-45d8-8037-e6695cd588c2
# ╠═816e841b-3f3b-4b8c-8c4b-4cd93c685ab7
# ╠═923b9503-5a8e-4dab-984d-6ca07d625bb4
# ╠═9938d1b8-8f80-45b0-bdff-48cd74e78fa4
# ╠═4771733f-83a5-4c13-b0ea-cb54241f7e19
# ╠═d8847644-4d5b-4113-8870-e8c3dba7e1a2
# ╠═5463ac9d-9bf7-4363-abb9-d1de01a88681
# ╠═39871429-722a-4745-9a67-62c2961fb268
# ╠═31a016af-82fa-41a1-b2e8-8adea326c11c
# ╠═31cfab8a-e219-484c-9428-1fef8ea16ca5
# ╠═c64a1c15-824a-49d2-ac0f-10573a1472c5
# ╠═8c9e6bd3-c548-4126-a6fb-a100e91f8747
# ╠═24425756-1e29-4e78-9398-6ef1820940b8
# ╠═4a74c5b1-d19b-4add-a0ca-3a239d8a7c75
# ╠═3914f4ec-4b70-4f05-9ec1-a9f70223bfd8
# ╠═94cf8ec3-5c81-4874-a331-aa0f15563808
# ╠═ffb3c3bb-827e-4173-90af-231bc45fef11
# ╠═2a1d988a-21fd-43b1-9701-dcfcf8f3d09b
# ╠═f51dc362-e8ab-494b-94f9-7ec63262afe4
# ╠═9ae7ea6f-6f70-4d92-9009-4ed3ccde3945
# ╠═5a0bd5e0-e3fb-47dc-8be2-fec23ea865a9
# ╠═2df2e495-5170-4db3-904b-0c272f69b182
# ╠═53142c10-19e9-42ba-b06d-85e326c8c575
# ╠═008b10fc-8e80-4783-ae37-1f18147d6f55
# ╠═c43d7521-0e3d-403c-83b6-d9cce7675931


using TraitSimulation
using GLM
using LinearAlgebra
using Random
using Test
using DataFrames
using VarianceComponentModels
using SnpArrays
using Suppressor
using DelimitedFiles

@testset "modelparameters" begin
	include("modelparameters.jl")
end

@testset "simulationmodels" begin
	include("simulationmodels.jl")
	include("matrixnormalsimulation.jl")
	include("writetopheno.jl")
end

@testset "simulatepower" begin
	include("vcm_power.jl")
	include("univariate_power.jl")
end


using SparseArrays
export choi2so, so2choi

"""
    permutedims_matrix(dims, p)

Returns a matrix `M` so that for any vector `v` of length `prod(dims)`,

    M*v == vec(permutedims(reshape(v, dims), p))

"""
function permutedims_matrix(dims, p)
    d = prod(dims)
    n = length(dims)
    sparse(reshape(PermutedDimsArray(reshape(I(d), (dims..., dims...)), (p..., (n + 1:2n)...)), (d, d)))
end

struct Choi2SOAtom <: Convex.AbstractExpr
    head::Symbol
    id_hash::UInt64
    children::Tuple{Convex.AbstractExpr}
    size::Tuple{Int,Int}
    dA::Int
    dB::Int
    function Choi2SOAtom(x::Convex.AbstractExpr, dA::Int, dB::Int)
        if dA^2 * dB^2 != length(x)
            error("Cannot Choi2SOAtom expression of size $(x.size) with dimensions ($(dA), $(dB))")
        end
        return new(:choi2so, objectid(x), (x,), (dB^2, dA^2), dA, dB)
    end
end

function Convex.sign(x::Choi2SOAtom)
    return Convex.sign(x.children[1])
end

function Convex.monotonicity(x::Choi2SOAtom)
    return (Convex.Nondecreasing(),)
end

function Convex.curvature(x::Choi2SOAtom)
    return Convex.ConstVexity()
end

function Convex.evaluate(x::Choi2SOAtom)
    J = Convex.evaluate(x.children[1])
    return choi2so(J, x.dA, x.dB)
end

function Convex.conic_form!(x::Choi2SOAtom, unique_conic_forms::Convex.UniqueConicForms = Convex.UniqueConicForms())
    if !Convex.has_conic_form(unique_conic_forms, x)
        objective = Convex.conic_form!(x.children[1], unique_conic_forms)
        dA = x.dA; dB = x.dB;
        M = permutedims_matrix((dA, dB, dA, dB), (2, 4, 1, 3))
        objective = M * objective
        Convex.cache_conic_form!(unique_conic_forms, x, objective)
    end

    return Convex.get_conic_form(unique_conic_forms, x)
end

"""
    choi2so(x::Convex.AbstractExpr, dA::Int, dB::Int)

Convert from an (unnormalized) Choi matrix of a quantum channel `T` to a matrix representation.
It is assumed `T` is a quantum channel from `A` to `B`, where `A` is `dA`-dimensional space and `B` a `dB`-dimensional space.
"""
choi2so(x::Convex.AbstractExpr, dA::Int, dB::Int) = Choi2SOAtom(x, dA, dB)
choi2so(J, dA, dB) = reshape(PermutedDimsArray(reshape(J, (dA, dB, dA, dB)), (2, 4, 1, 3)), (dB^2, dA^2))

struct SO2ChoiAtom <: Convex.AbstractExpr
    head::Symbol
    id_hash::UInt64
    children::Tuple{Convex.AbstractExpr}
    size::Tuple{Int,Int}
    dA::Int
    dB::Int
    function SO2ChoiAtom(x::Convex.AbstractExpr, dA::Int, dB::Int)
        if !(dA^2 * dB^2 == length(x) && size(x) == (dB^2, dA^2))
            error("Cannot Choi2SOAtom expression of size $(x.size) with dimensions ($(dA), $(dB))")
        end
        return new(:so2choi, objectid(x), (x,), (dA*dB, dA*dB), dA, dB)
    end
end

function Convex.sign(x::SO2ChoiAtom)
    return Convex.sign(x.children[1])
end

function Convex.monotonicity(x::SO2ChoiAtom)
    return (Convex.Nondecreasing(),)
end

function Convex.curvature(x::SO2ChoiAtom)
    return Convex.ConstVexity()
end

function Convex.evaluate(x::SO2ChoiAtom)
    T = Convex.evaluate(x.children[1])
    return so2choi(T, x.dA, x.dB)
end

function Convex.conic_form!(x::SO2ChoiAtom, unique_conic_forms::Convex.UniqueConicForms = Convex.UniqueConicForms())
    if !Convex.has_conic_form(unique_conic_forms, x)
        objective = Convex.conic_form!(x.children[1], unique_conic_forms)
        dA = x.dA; dB = x.dB;
        M = permutedims_matrix((dB, dB, dA, dA), (3, 1, 4, 2))
        objective = M * objective
        Convex.cache_conic_form!(unique_conic_forms, x, objective)
    end

    return Convex.get_conic_form(unique_conic_forms, x)
end

"""
    so2choi(x::Convex.AbstractExpr, dA::Int, dB::Int)

Convert from a matrix representation of a quantum channel `T` to the Choi representation.
It is assumed `T` is a quantum channel from `A` to `B`, where `A` is `dA`-dimensional space and `B` a `dB`-dimensional space.
"""
so2choi(x::Convex.AbstractExpr, dA::Int, dB::Int) = SO2Choi(x, dA, dB)
so2choi(T, dA, dB) = reshape(PermutedDimsArray(reshape(T, (dB, dB, dA, dA)), (3, 1, 4, 2)), (dA*dB, dA*dB))


push!(LOAD_PATH, joinpath(@__DIR__, ".."))

using BenchmarkTools
using CUDA
using DataDeps
using Oceananigans
using Benchmarks
using Statistics
using Oceananigans.Solvers: sparse_approximate_inverse

function benchmark_spai_preconditioner(N, ε, nzrel, inverse)

    grid = RectilinearGrid(CPU(), size=(N, N, 1), extent=(1, 1, 1))

    model = HydrostaticFreeSurfaceModel(
                      grid = grid,
        momentum_advection = VectorInvariant(),
              free_surface = ImplicitFreeSurface(solver_method=:HeptadiagonalIterativeSolver, preconditioner_method = nothing)
    )

    # to correctly create the matrix
    time_step!(model, 1)

    matrix = model.free_surface.implicit_step_solver.matrix_iterative_solver.matrix

    if !inverse
        trial = @benchmark begin
            CUDA.@sync blocking = true sparse_approximate_inverse($matrix, ε = $ε, nzrel = $nzrel)
        end samples = 5
    else
        trial = @benchmark begin
            CUDA.@sync blocking = true inv(Array($matrix))
        end samples = 5
    end

    return trial
end

N       = [64, 128, 256]
ε       = [0.1, 0.3]
nzrel   = [0.5, 1.0, 2.0]
inverse = (false, )

# Run and summarize benchmarks
print_system_info()
suite = run_benchmarks(benchmark_spai_preconditioner; N, ε, nzrel, inverse)

df = benchmarks_dataframe(suite)
benchmarks_pretty_table(df, title="SPAI preconditioner benchmarks")