text:
##### Diffusion approximation for Slab Geometry #####

#=
@article{contini1997photon,
  title={Photon migration through a turbid slab described by a model based on diffusion approximation. I. Theory},
  author={Contini, Daniele and Martelli, Fabrizio and Zaccanti, Giovanni},
  journal={Applied optics},
  volume={36},
  number={19},
  pages={4587--4599},
  year={1997},
  publisher={Optical Society of America}
}
@article{kienle1997improved,
  title={Improved solutions of the steady-state and the time-resolved diffusion equations for reflectance from a semi-infinite turbid medium},
  author={Kienle, Alwin and Patterson, Michael S},
  journal={JOSA A},
  volume={14},
  number={1},
  pages={246--254},
  year={1997},
  publisher={Optical Society of America}
}
@inproceedings{martelli2009light,
  title={Light propagation through biological tissue and other diffusive media: theory, solutions, and software},
  author={Martelli, Fabrizio and Del Bianco, Samuele and Ismaelli, Andrea},
  year={2009},
  organization={Society of Photo-Optical Instrumentation Engineers}
}
=#

### Fluence TD ###
"""
    fluence_DA_slab_TD(t, β::Array{Float64,1}, ρ::Float64, ndet::Float64, nmed::Float64, s::Float64, z::Float64)

Compute the time-domain fluence from a slab geometry (x, y->inf, z-> finite) with Eqn. 33 Contini 97.

# Arguments
- `t`: the time vector (ns).
- `β::Array{Float64,1}`: the optical properties [μa, μs'] (cm⁻¹)
- `ρ::Float64`: the source detector separation (cm⁻¹)
- `ndet::Float64`: the boundary's index of refraction (air or detector)
- `nmed::Float64`: the sample medium's index of refraction
- `s::Float64`: the thickness (z-depth) of the slab (cm)
- `z::Float64`: the z-depth coordinate (cm)

# Examples
julia> fluence_DA_slab_TD(0:1:5, [0.1,10.0], 1.0, 1.0, 1.0, 2.0, 0.0)
"""
function fluence_DA_slab_TD(t, β::Array{Float64,1}, ρ::Float64, ndet::Float64, nmed::Float64, s::Float64, z::Float64)
	  n::Float64 = nmed/ndet
      μa::Float64 = β[1]
      μsp::Float64 = β[2]
      D::Float64 = 1/3μsp
      ν::Float64 = 29.9792345/nmed
      xs::UnitRange{Int64} = -10:10

      Rt1 = Array{Float64}(undef, length(t))
	  Rt2 = zeros(Float64, length(t))

		@inbounds begin

			if n == 1.0
				A= 1.0
			elseif n > 1.0
				A = 504.332889 - 2641.00214n + 5923.699064n^2 - 7376.355814n^3 +
		 		5507.53041n^4 - 2463.357945n^5 + 610.956547n^6 - 64.8047n^7
			else
				A = 3.084635 - 6.531194n + 8.357854n^2 - 5.082751n^3
			end

			zs::Float64 = 1/μsp
			ze::Float64 = 2A*D

                Threads.@threads for n in eachindex(t)
                    Rt1[n] = @. ν*exp(-(ρ^2/(4D*ν*t[n])) - μa*ν*t[n])/((4π*D*ν*t[n])^(3/2))

					for m in xs

						zmp = 2m.*(s + 2*ze) .+ zs
						zmm = 2m.*(s + 2*ze) .- 2*ze .- zs

			 			Rt2[n] += exp(-((z - zmp)^2 / (4D*ν*t[n]))) - exp(-((z - zmm)^2 / (4D*ν*t[n])))
                    end

                    Rt1[n] *= Rt2[n]

                    if isnan(Rt1[n])
                        Rt1[n] = 0
                    end

				end
		end

	return Rt1
end

### Reflectance TD ###
"""
    refl_DA_slab_TD(t, β::Array{Float64,1}, ρ::Float64,ndet::Float64, nmed::Float64, s::Float64)

Compute the time-domain reflectance from a slab geometry (x,y->inf, z-> finite).

# Arguments
- `t`: the time vector (ns).
- `β::Array{Float64,1}`: the optical properties [μa, μs'] (cm⁻¹)
- `ρ::Float64`: the source detector separation (cm⁻¹)
- `ndet::Float64`: the boundary's index of refraction (air or detector)
- `nmed::Float64`: the sample medium's index of refraction
- `s::Float64`: the thickness (z-depth) of the slab (cm)

# Examples
julia> refl_DA_slab_TD(0:1:5, [0.1,10.0], 1.0, 1.0, 1.0, 2.0)
"""
function refl_DA_slab_TD(t, β::Array{Float64,1}, ρ::Float64,ndet::Float64, nmed::Float64, s::Float64)
	  n::Float64 = nmed/ndet
      μa::Float64 = β[1]
      μsp::Float64 = β[2]
      D::Float64 = 1/3μsp
      ν::Float64 = 29.9792345/nmed
      xs::UnitRange{Int64} = -10:10

      Rt1 = Array{Float64}(undef, length(t))
	  Rt2 = zeros(Float64, length(t))

		@inbounds begin

			if n == 1.0
				A= 1.0
			elseif n > 1.0
				A = 504.332889 - 2641.00214n + 5923.699064n^2 - 7376.355814n^3 +
		 		5507.53041n^4 - 2463.357945n^5 + 610.956547n^6 - 64.8047n^7
			else
				A = 3.084635 - 6.531194n + 8.357854n^2 - 5.082751n^3
			end

			zs::Float64 = 1/μsp
			ze::Float64 = 2A*D

                Threads.@threads for n in eachindex(t)
                    Rt1[n] = @. -exp(-(ρ^2/(4D*ν*t[n])) - μa*ν*t[n])/(2*(4π*D*ν)^(3/2)*t[n]^(5/2))

					for m in xs

						z3m = -2m*s - 4m*ze - zs
						z4m = -2m*s - (4m - 2)*ze + zs

			 			Rt2[n] += z3m*exp(-(z3m^2 / (4D*ν*t[n]))) - z4m*exp(-(z4m^2 / (4D*ν*t[n])))
					end

					Rt1[n] *= Rt2[n]

                    if isnan(Rt1[n])
                        Rt1[n] = 0
                    end
				end
			end

	return Rt1
end

### Transmittance TD ###
"""
    trans_DA_slab_TD(t, β::Array{Float64,1}, ρ::Float64,ndet::Float64, nmed::Float64, s::Float64)

Compute the time-domain transmittance from a slab geometry (x,y->inf, z-> finite) with Eqn. 39 from Contini 97.

# Arguments
- `t`: the time vector (ns).
- `β::Array{Float64,1}`: the optical properties [μa, μs'] (cm⁻¹)
- `ρ::Float64`: the source detector separation (cm⁻¹)
- `ndet::Float64`: the boundary's index of refraction (air or detector)
- `nmed::Float64`: the sample medium's index of refraction
- `s::Float64`: the thickness (z-depth) of the slab (cm)

# Examples
julia> trans_DA_slab_TD(0:1:5, [0.1,10.0], 1.0, 1.0, 1.0, 2.0)
"""
function trans_DA_slab_TD(t, β::Array{Float64,1}, ρ::Float64,ndet::Float64, nmed::Float64, s::Float64)
	  n::Float64 = nmed/ndet
      μa::Float64 = β[1]
      μsp::Float64 = β[2]
      D::Float64 = 1/3μsp
      ν::Float64 = 29.9792345/nmed
      xs::UnitRange{Int64} = -10:10

      Rt1 = Array{Float64}(undef, length(t))
	  Rt2 = zeros(Float64, length(t))

		@inbounds begin

			if n == 1.0
				A= 1.0
			elseif n > 1.0
				A = 504.332889 - 2641.00214n + 5923.699064n^2 - 7376.355814n^3 +
		 		5507.53041n^4 - 2463.357945n^5 + 610.956547n^6 - 64.8047n^7
			else
				A = 3.084635 - 6.531194n + 8.357854n^2 - 5.082751n^3
			end

			zs::Float64 = 1/μsp
			ze::Float64 = 2A*D

                Threads.@threads for n in eachindex(t)
                    Rt1[n] = @. exp(-(ρ^2/(4D*ν*t[n])) - μa*ν*t[n])/(2*(4π*D*ν)^(3/2)*t[n]^(5/2))

					for m in xs

						z1m = (1 - 2*m)*s - 4m*ze - zs
						z2m = (1 - 2*m)*s - (4m - 2)*ze + zs

			 			Rt2[n] += z1m*exp(-(z1m^2 / (4D*ν*t[n]))) - z2m*exp(-(z2m^2 / (4D*ν*t[n])))
					end

					Rt1[n] *= Rt2[n]

                    if isnan(Rt1[n])
                        Rt1[n] = 0
                    end
				end
			end

	return Rt1
end

### Fluence CW ###
"""
    fluence_DA_slab_CW(ρ::Float64, β::Array{Float64,1}, ndet::Float64, nmed::Float64, s::Float64, z::Float64)

Compute the steady-state fluence from a slab geometry (x,y->inf, z-> finite).

# Arguments
- `ρ::Float64`: the source detector separation (cm⁻¹)
- `β::Array{Float64,1}`: the optical properties [μa, μs'] (cm⁻¹)
- `ndet::Float64`: the boundary's index of refraction (air or detector)
- `nmed::Float64`: the sample medium's index of refraction
- `s::Float64`: the thickness (z-depth) of the slab (cm)
- `z::Float64`: the z-depth within slab (cm)

# Examples
julia> fluence_DA_slab_CW(1.0, [0.1,10.], 1.0,1.0, 2.0, 0.)
"""
function fluence_DA_slab_CW(ρ::Float64, β::Array{Float64,1}, ndet::Float64, nmed::Float64, s::Float64, z::Float64)

      n::Float64 = nmed/ndet
      μa::Float64 = β[1]
      μsp::Float64 = β[2]
      D::Float64 = 1/3μsp
      μeff::Float64 = sqrt(3*μa*μsp)
      xs::UnitRange{Int64} = -10:10

	  ϕ = 0.0

		if n == 1.0
			A= 1.0
		elseif n > 1.0
			A = 504.332889 - 2641.00214n + 5923.699064n^2 - 7376.355814n^3 +
			5507.53041n^4 - 2463.357945n^5 + 610.956547n^6 - 64.8047n^7
		else
			A = 3.084635 - 6.531194n + 8.357854n^2 - 5.082751n^3
		end

		zs::Float64 = 1/μsp
		ze::Float64 = 2A*D

        Threads.@threads for m in xs

			zmp = 2m.*(s + 2*ze) .+ zs
			zmm = 2m.*(s + 2*ze) .- 2*ze .- zs

            ϕ += exp(-μeff*sqrt(ρ^2 + (z - zmp)^2))/(sqrt(ρ^2 + (z - zmp)^2))
            ϕ -= exp(-μeff*sqrt(ρ^2 + (z - zmm)^2))/(sqrt(ρ^2 + (z - zmm)^2))
        end

        if isnan(ϕ)
            ϕ = 0
        end

	return ϕ/(4*π*D)
end

max_stars_repo_name: sammzer/LightPropagation

max_issues_repo_name: sammzer/LightPropagation

max_forks_repo_name: sammzer/LightPropagation


text:
function FieldTransformGenerator(arg0::Field, arg1::jint, arg2::TransformProvider, arg3::jdouble)
    return FieldTransformGenerator((Field, jint, TransformProvider, jdouble), arg0, arg1, arg2, arg3)
end

function generate(obj::FieldTransformGenerator, arg0::AbsoluteDate, arg1::AbsoluteDate)
    return jcall(obj, "generate", List, (AbsoluteDate, AbsoluteDate), arg0, arg1)
end

ext: jl


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

using Markdown
using InteractiveUtils

# ╔═╡ 4de51e08-3854-11eb-10f2-d17580deef9c
begin
	using Pkg
	Pkg.add(["AbstractTrees"])
end

# ╔═╡ 2ac17898-3856-11eb-1235-4fd15bc1592c
begin
	using AbstractTrees
end

# ╔═╡ 3e0c149a-3350-11eb-0b17-7bea5d1a28d4
begin
	DAY = 7
	md"# 2020 Day $DAY"
end

# ╔═╡ 18080fe2-339b-11eb-34f0-f393638506fc
md"## Part 1"

# ╔═╡ 4e084a4e-339b-11eb-3ffc-41118215f226
md"*Bags*"

# ╔═╡ 24e1613c-339b-11eb-2976-59884053d1ee
input = read("../data/day_$DAY.txt", String)

# ╔═╡ df1d1be2-36c6-11eb-07c6-493a612dd182
test_input = """
light red bags contain 1 bright white bag, 2 muted yellow bags.
dark orange bags contain 3 bright white bags, 4 muted yellow bags.
bright white bags contain 1 shiny gold bag.
muted yellow bags contain 2 shiny gold bags, 9 faded blue bags.
shiny gold bags contain 1 dark olive bag, 2 vibrant plum bags.
dark olive bags contain 3 faded blue bags, 4 dotted black bags.
vibrant plum bags contain 5 faded blue bags, 6 dotted black bags.
faded blue bags contain no other bags.
dotted black bags contain no other bags.
"""

# ╔═╡ fefa6924-3850-11eb-067a-5f62b6570b06

# ╔═╡ 0f76e99c-3851-11eb-113c-0136765cd4c4

# ╔═╡ 5f974aa2-3851-11eb-2198-09a8c6d7881a
function get_edges(line::AbstractString)
	raw_container, contents = split(line,"contain")
	container = raw_container[1:end-6]
	bags = (
		split(contents,",")
		.|> strip
		.|> s -> replace(s, "s." => "")
		# container, content, number
		.|> s -> [container, strip(s[3:end-4]), parse(Int,s[1])]
		)
end

# ╔═╡ 3858684c-3857-11eb-3506-4f1c87a22921
function input2edges(input)
	(
	[s for s in split(strip(input),"\n") if !occursin("no other bags",s)]
	.|> get_edges
	|> nested_list -> [item for sublist in nested_list for item in sublist]
)
end

# ╔═╡ d7a2264a-3859-11eb-3ac9-a176933f2063
function edges2colors(edges)
	container, content, number = zip(edges...)
	union(Set(container),Set(content))
end

# ╔═╡ 00aec94e-385a-11eb-3985-b9a159435b1a
function edges2mapping(edges)
	colors = edges2colors(edges)
	d = Dict(c => [] for c in colors)
	for (o, i, n) in edges
		push!(d[i],o)
	end
	return d
end

# ╔═╡ 46c193da-385a-11eb-31a6-25e4eda23b35

# ╔═╡ 4fef820e-385b-11eb-1f9a-ff3310d2340c
function find_parent!(outer_colors,d,c)
	for color in d[c]
		if color ∉ outer_colors
			push!(outer_colors, color)
			find_parent!(outer_colors,d,color)
		end
	end
end

# ╔═╡ 31101e58-385f-11eb-2afa-55f88de38916
e = input2edges(input)

# ╔═╡ 6310ea54-385f-11eb-25b2-a353f8585388
mapping = edges2mapping(e)

# ╔═╡ 863e77da-385f-11eb-19a6-53f6116e4a24
colors = [i for i in edges2colors(e)]

# ╔═╡ 7e22f556-3893-11eb-38d2-1bfc0441a8ad
function find_parents(e,mapping,color)
	outer_colors = Set()
	return find_parent!(outer_colors,mapping,color)
end

# ╔═╡ 91204704-38bf-11eb-298c-6995aa085a06
begin
	outer_colors = Set()
	find_parent!(outer_colors,mapping,"shiny gold")
	outer_colors
end

# ╔═╡ 1b524b64-385e-11eb-17b1-d920e32aab0c

# ╔═╡ 2801d0ec-38c1-11eb-28dc-3ff0a89ee9a7
e

# ╔═╡ 5f797e22-339d-11eb-1ff3-dda90f946838
md"## Part 2"

# ╔═╡ e3177eae-339e-11eb-3f9f-1f7230976832
md"*<Insert Description>*"

# ╔═╡ Cell order:
# ╠═3e0c149a-3350-11eb-0b17-7bea5d1a28d4
# ╟─18080fe2-339b-11eb-34f0-f393638506fc
# ╠═4e084a4e-339b-11eb-3ffc-41118215f226
# ╠═4de51e08-3854-11eb-10f2-d17580deef9c
# ╠═2ac17898-3856-11eb-1235-4fd15bc1592c
# ╠═24e1613c-339b-11eb-2976-59884053d1ee
# ╠═df1d1be2-36c6-11eb-07c6-493a612dd182
# ╟─fefa6924-3850-11eb-067a-5f62b6570b06
# ╟─0f76e99c-3851-11eb-113c-0136765cd4c4
# ╠═5f974aa2-3851-11eb-2198-09a8c6d7881a
# ╠═3858684c-3857-11eb-3506-4f1c87a22921
# ╠═d7a2264a-3859-11eb-3ac9-a176933f2063
# ╠═00aec94e-385a-11eb-3985-b9a159435b1a
# ╟─46c193da-385a-11eb-31a6-25e4eda23b35
# ╠═4fef820e-385b-11eb-1f9a-ff3310d2340c
# ╠═31101e58-385f-11eb-2afa-55f88de38916
# ╠═6310ea54-385f-11eb-25b2-a353f8585388
# ╠═863e77da-385f-11eb-19a6-53f6116e4a24
# ╠═7e22f556-3893-11eb-38d2-1bfc0441a8ad
# ╠═91204704-38bf-11eb-298c-6995aa085a06
# ╟─1b524b64-385e-11eb-17b1-d920e32aab0c
# ╠═2801d0ec-38c1-11eb-28dc-3ff0a89ee9a7
# ╠═5f797e22-339d-11eb-1ff3-dda90f946838
# ╟─e3177eae-339e-11eb-3f9f-1f7230976832

max_stars_repo_name: RensDimmendaal/advent-of-code-julia

max_issues_repo_name: RensDimmendaal/advent-of-code-julia

max_forks_repo_name: RensDimmendaal/advent-of-code-julia


text:
export partial_trace, trace_distance, post_CPTP!, max_ent_state!, iden, bases_create!, full_process

"""
    post_CPTP!(ρ)

Applies a CPTP contraction to the input Hermitian operator
Input arguments:
    ρ: Hermitian operator to be transformed to a positive semi-definite version
"""
function post_CPTP!(rho::Array{T,2}) where {T <: Number}

    A = eigen(rho)
    A.values[real(A.values) .<0.] .= 0.
    rho[:,:] = A.vectors*Diagonal(A.values)*A.vectors'
    nothing
end

"""
    trace_distance(A, B)

Calculates the trace distance between two matrices
Input arguments:
    A, B: matrices of the same dimensions
Output arguments:
    : trace distance
"""
function trace_distance(A::Array{T1,2},B::Array{T2,2}) where {T1 <: Number, T2 <: Number}

    tr(sqrt((A-B)'*(A-B)))/2
end

"""
    partial_trace(rho, qs, out, dim, num)

This function calculates the partial trace of a density matrix
Input arguments:
    rho: full density matrix
    qs: a vector of Ints of length num, with a 1 if the qudit is to be traced out and a 0 otherwise
    dim: vector of Ints of the dimensions of the each qudit (same order as qs)
    num: total number of subsystems
Output:
       rhoPT: the partial trace density matrix
"""
function partial_trace(rho::Array{T,2},qs::Array{Int64,1},dim::Array{Int64,1},num::Int64) where {T <: Number}

    rhoPT = rho
    for ii = 1:1:num
        tempV = zeros(dim[ii])

        if qs[ii] == 1
            dim1 = dim[1:ii]
            qs1 = qs[1:ii]
            tdim = div(prod(dim),prod(dim1[qs1.==1]))
            temp = zeros(tdim,tdim)

            for jj = 1:1:dim[ii]
                dimt = dim[1:(ii-1)]
                qst = qs[1:(ii-1)]
                tdim2 = prod(dimt[qst.==0])
                tdim3 = prod(dim[(ii+1):end])

                unitV(tempV,jj)
                tempOp = kron(Matrix{ComplexF64}(I,tdim2,tdim2),kron(tempV,Matrix{ComplexF64}(I,tdim3,tdim3)))
                temp = temp + tempOp'*rhoPT*tempOp
            end
            rhoPT = temp
        end
    end
    return rhoPT
end

"""In place version of the `partial_trace` function. Note the explicit type required for out."""
function partial_trace(out::Array{ComplexF64,2},rho::Array{T,2},qs::Array{Int64,1},dim::Array{Int64,1},num::Int64) where {T <: Number}

    rhoPT = rho
    for ii = 1:1:num
        tempV = zeros(dim[ii])

        if qs[ii] == 1
            dim1 = dim[1:ii]
            qs1 = qs[1:ii]
            tdim = div(prod(dim),prod(dim1[qs1.==1]))
            temp = zeros(tdim,tdim)

            for jj = 1:1:dim[ii]
                dimt = dim[1:(ii-1)]
                qst = qs[1:(ii-1)]
                tdim2 = prod(dimt[qst.==0])
                tdim3 = prod(dim[(ii+1):end])

                unitV(tempV,jj)
                tempOp = kron(Matrix{ComplexF64}(I,tdim2,tdim2),kron(tempV,Matrix{ComplexF64}(I,tdim3,tdim3)))
                temp = temp + tempOp'*rhoPT*tempOp
            end
            rhoPT = temp
        end
    end
    out[:,:] = rhoPT
    nothing
end

"""
    unitV(tempV, el)

This function creates a unit vector
Input arguments:
    tempV: vector of zeros that will be overwritten by the unit vector
    el: location of the element to set to 1.
"""
function unitV(tempV::Array{Float64,1},el::Int64)
    fill!(tempV,0.)
    tempV[el] = 1.
    nothing
end

"""
    full_process(out, ρ, χ, I)

This function calculates the output state for a given input state and set of two-qubit quantum processes
Input:
    out: output state
    rho: input state
    Choi: [Nprocesses, dim^2, dim^2]cell array of Choi states describing the processes
    iden: identity the same size as rho
"""
function full_process(out::Array{ComplexF64,2},rho::Array{ComplexF64,2},Choi::Array{ComplexF64,3},iden::Array{ComplexF64,2})

    out[:,:] = rho
    for ii = 1:1:size(Choi,1)
        out[:,:] = partial_trace(kron(transpose(out),iden)*Choi[ii,:,:],[1,0],[size(rho,1),size(rho,1)],2).*size(rho,1)
    end
end

"""
    iden(dim)

This function creates an identity matrix of complex floats
Input:
    dim: dimension of the identity matrix
Output:
    identity matrix of complex floats of dimension dim
"""
function iden(dim::Int)
    return Matrix{ComplexF64}(I,dim,dim)
end

"""
    bases_create!(Bases, N, Npairs, pair_order)

This function creates a two-qubit operator basis dilated to N qubits, for each pair of qubits in the full set
Input:
    Bases: [Npairs,16,2^Nqubits,2^Nqubits] size array that stores the bases
    N: number of qubits
    Npairs: number of qubit pairs
    pair_order: ordering of the qubit pairs, qubit number counting starts at 1
Output:
    nothing (modifies the elements of Bases)
"""
function bases_create!(Bases::Array{T,4},N::Int,Npairs::Int,pair_order::Array{Int,2}) where {T <: Number}

    Bases[:,:,:,:] = convert(Array{ComplexF64,4},Bases)

    tempV1 = zeros(2)
    tempV2 = zeros(2)

    for nn = 1:1:Npairs
        for aa = 1:1:2
            for bb = 1:1:2
                for cc = 1:1:2
                    for dd = 1:1:2

                        unitV(tempV1,aa)
                        unitV(tempV2,bb)
                        tempab = kron(iden(2^(pair_order[nn,1]-1)),kron(tempV1,kron(iden(2^(pair_order[nn,2]-pair_order[nn,1]-1)),kron(tempV2,iden(2^(N-pair_order[nn,2]))))))

                        unitV(tempV1,cc)
                        unitV(tempV2,dd)
                        tempcd = kron(iden(2^(pair_order[nn,1]-1)),kron(tempV1,kron(iden(2^(pair_order[nn,2]-pair_order[nn,1]-1)),kron(tempV2,iden(2^(N-pair_order[nn,2]))))));

                        Bases[nn,4*((2*(aa-1)+bb)-1)+2*(cc-1)+dd,:,:] = tempab*tempcd'
                    end
                end
            end
        end
    end
    nothing
end

"""
    max_ent_state!(dim, mvec, mrho)

This function creates a maximally entangled state. It can return either a state vector or density matrix (or both). Note the syntax used!
Input:
    dim: Hilbert space dimension of the system
Keyword optional input:
    mvec: empty vector to store the maximally entangled state vector
    mrho: empty matrix to store the maximally entangled density matrix
Output:
    nothing
"""
function max_ent_state!(dim::Int;mvec::Vector{Float64}=zeros(dim*dim),mrho::Array{Float64,2}=zeros(dim*dim,dim*dim))

    tempV = zeros(dim)
    for mm = 1:1:dim
        unitV(tempV,mm)
        mvec[:] = mvec[:] + kron(tempV,tempV)
    end
    mvec[:] = mvec[:]./sqrt(dim)

    mrho[:,:] = mvec*mvec'
    nothing
end


text:
#
# show() methods to make output of complex structures readable
#

# printstyled(IOContext(io, :color => true), "string", color=:red)

import Base: show

spaces = "  "

function _indent_level!(io::IO, delta::Int)
    level = get(io, :indent_level, 0)
    return IOContext(io, :indent_level => max(level + delta, 0))
end

indent(io::IO) = _indent_level!(io, 1)

"""
    printable(value)

Return a value that is not nothing.
"""
printable(value) = (value === nothing ? ":nothing:" : value)

function print_indented(io::IO, args...)
    level = get(io, :indent_level, 0)
    print(io, repeat(spaces, level), args...)
    nothing
end

function _show_field(io::IO, name::Symbol, value; show_empty=true)
    if !show_empty && isempty(value)
        return
    end

    print(io, "\n")
    print_indented(io, name, ": ")
    show(io, value)
end

function _show_field(io::IO, name::Symbol, dict::AbstractDict; show_empty=true)
    if !show_empty && isempty(dict)
        return
    end

    print(io, "\n")
    print_indented(io, name, ": ", typeof(dict))
    io = indent(io)
    for (k, v) in dict
        print(io, "\n")
        print_indented(io, k, " => ")
        show(io, v)
    end
    nothing
end

function _show_field(io::IO, name::Symbol, vec::Vector{T}; show_empty=true) where T
    print(io, "\n")
    print_indented(io, name, ": ", typeof(vec))

    count = length(vec)
    elide = false
    max_shown = 5
    if count > max_shown
        last = vec[end]
        vec = vec[1:max_shown-1]
        elide = true
    end

    for (i, value) in enumerate(vec)
        print(io, "\n")
        print_indented(io, "$i: ", value)
    end

    if elide
        print(io, "\n")
        print_indented(io, "...\n")
        print_indented(io, "$count: ")
        show(io, last)
    end
    nothing
end

function _show_field(io::IO, name::Symbol, vec::Vector{<: AbstractMimiType}; show_empty=true)
    if !show_empty && isempty(vec)
        return
    end

    print(io, "\n")
    print_indented(io, name, ": ")
    io = indent(io)
    for (i, value) in enumerate(vec)
        print(io, "\n")
        print_indented(io, "$i: ", value)
    end
end

function _show_fields(io::IO, obj, names; show_empty=true)
    for name in names
        value = getfield(obj, name)
        _show_field(io, name, value)
    end
    nothing
end

function _show_datum_def(io::IO, obj::AbstractDatumDef)
    print(io, typeof(obj), "($(obj.name)::$(obj.datatype))")
    io = indent(io)
    _show_field(io, :comp_path, obj.comp_path)
    _show_field(io, :dim_names, obj.dim_names)

    for field in (:description, :unit)
        _show_field(io, field, getfield(obj, field), show_empty=false)
    end
end

function show(io::IO, obj::ComponentId)
    print(io, "ComponentId($(obj.module_obj).$(obj.comp_name))")
    nothing
end

function show(io::IO, obj::ComponentPath)
    print(io, "ComponentPath$(obj.names)")
    nothing
end

function show(io::IO, obj::AbstractDimension)
    print(io, keys(obj))
    nothing
end

show(io::IO, obj::VariableDef) = _show_datum_def(io, obj)

function show(io::IO, obj::ParameterDef)
    _show_datum_def(io, obj)
    _show_field(indent(io), :default, obj.default)
end

function show(io::IO, obj::AbstractMimiType)
    print(io, typeof(obj))

    # If a name is present, print it as type(name)
    fields = fieldnames(typeof(obj))
    pos = findfirst(x -> x == :name, fields)

    if pos !== nothing
        print(io, "($(obj.name))")
        fields = deleteat!([fields...], pos)
    end

    _show_fields(indent(io), obj, fields)
end

function show(io::IO, obj::AbstractComponentDef)
    print(io, nameof(typeof(obj)), " id:", objectid(obj))

    fields = fieldnames(typeof(obj))

    # skip the 'namespace' field since it's redundant
    fields = [f for f in fields if f != :namespace]

    # Don't print parent or root since these create circular references
    for field in (:parent, :root)
        pos = findfirst(x -> x == field, fields)
        if pos !== nothing
            value = getfield(obj, field)
            name = printable(value === nothing ? nothing : nameof(value))
            fields = deleteat!([fields...], pos)
            print(io, "\n")
            print_indented(indent(io), "$field: $(typeof(value))($name)")
        end
    end

    io = indent(io)
    _show_fields(io, obj, fields)

    if obj isa AbstractComponentDef
        # print an abbreviated namespace
        print(io, "\n")
        print_indented(io, "namespace:")
        io = indent(io)
        for (name, item) in obj.namespace
            print(io, "\n")
            print_indented(io, name, ": ")
            show(io, typeof(item))
        end
    end
end

function show(io::IO, obj::AbstractDatumReference)
    print(io, nameof(typeof(obj)), "(name=:$(obj.name) path=$(obj.comp_path))")
end

function show(io::IO, obj::ModelInstance)
    # Don't print full type signature since it's shown in .variables and .parameters
    print(io, "ModelInstance")

    # Don't print the mi's ModelDef since it's redundant
    fields = fieldnames(typeof(obj))
    pos = findfirst(x -> x == :md, fields)
    fields = deleteat!([fields...], pos)

    io = indent(io)
    _show_fields(io, obj, fields)
    print(io, "\n")
    print_indented(io, "md: (not shown)")
end

#
# We might want to use a version of this to simplify. Imported forward from Mimi v0.9.1.
#
function _show(io::IO, obj::Model, which::Symbol)

    println(io, "Mimi.Model")
    md = obj.md
    mi = obj.mi

    println(io, "  Module: $(md.comp_id.module_obj)")

    println(io, "  Components:")
    for comp in values(components(md))
        println(io, "    $(comp.comp_id)")
    end

    if which == :full
        println(io, "  Dimensions:")
        for (k, v) in md.dim_dict
            println(io, "    $k => $v")
        end

        println(io, "  Internal Connections:")
        for conn in md.internal_param_conns
            println(io, "    $(conn)")
        end

        println(io, "  External Connections:")
        for conn in md.external_param_conns
            println(io, "    $(conn)")
        end

        println(io, "  Backups: $(md.backups)")
        println(io, "  Number type: $(md.number_type)")
    end
    println(io, "  Built: $(mi !== nothing)")
end

Base.show(io::IO, obj::Model) = _show(io, obj, :full)

Base.show(io::IO, ::MIME"text/plain", obj::Model) = _show(io, obj, :short)


text:
using Test
using ArnoldiMethod: OrderBy, OrderReverse, OrderPerm
using Base.Order: lt

@testset "Stable permutation ordering" begin
    xs = [1+3im, 1-3im, 4]
    @testset "Forward with f = $f" for f in (real, abs)
        ord = OrderPerm(xs, OrderBy(f))

        # Test whether the order is stable!
        @test lt(ord, 1, 2)
        @test !lt(ord, 2, 1)
        @test lt(ord, 1, 3)
        @test lt(ord, 2, 3)
        @test sort!([1,2,3], QuickSort, ord) == [1, 2, 3]
    end

    @testset "Backward with f = $f" for f in (real, abs)
        ord = OrderPerm(xs, OrderReverse(OrderBy(f)))

        # Test whether the order is stable even backwards -- note that i
        @test lt(ord, 1, 2)
        @test !lt(ord, 2, 1)
        @test lt(ord, 3, 1)
        @test lt(ord, 3, 2)
        @test sort!([1,2,3], QuickSort, ord) == [3, 1, 2]
    end
end


text:
_nrow(df::AbstractDataFrame) = nrow(df)
_nrow(x::NamedTuple{<:Any, <:Tuple{Vararg{AbstractVector}}}) =
    isempty(x) ? 0 : length(x[1])
_ncol(df::AbstractDataFrame) = ncol(df)
_ncol(x::Union{NamedTuple, DataFrameRow}) = length(x)

function _combine_multicol(firstres, fun::Base.Callable, gd::GroupedDataFrame,
                           incols::Union{Nothing, AbstractVector, Tuple, NamedTuple})
    firstmulticol = firstres isa MULTI_COLS_TYPE
    if !(firstres isa Union{AbstractVecOrMat, AbstractDataFrame,
                            NamedTuple{<:Any, <:Tuple{Vararg{AbstractVector}}}})
        idx_agg = Vector{Int}(undef, length(gd))
        fillfirst!(nothing, idx_agg, 1:length(gd.groups), gd)
    else
        idx_agg = nothing
    end
    return _combine_with_first(wrap(firstres), fun, gd, incols,
                               Val(firstmulticol), idx_agg)
end

function _combine_with_first(first::Union{NamedTuple, DataFrameRow, AbstractDataFrame},
                             f::Base.Callable, gd::GroupedDataFrame,
                             incols::Union{Nothing, AbstractVector, Tuple, NamedTuple},
                             firstmulticol::Val, idx_agg::Union{Nothing, AbstractVector{<:Integer}})
    extrude = false

    if first isa AbstractDataFrame
        n = 0
        eltys = eltype.(eachcol(first))
    elseif first isa NamedTuple{<:Any, <:Tuple{Vararg{AbstractVector}}}
        n = 0
        eltys = map(eltype, first)
    elseif first isa DataFrameRow
        n = length(gd)
        eltys = [eltype(parent(first)[!, i]) for i in parentcols(index(first))]
    elseif firstmulticol == Val(false) && first[1] isa Union{AbstractArray{<:Any, 0}, Ref}
        extrude = true
        first = wrap_row(first[1], firstmulticol)
        n = length(gd)
        eltys = (typeof(first[1]),)
    else # other NamedTuple giving a single row
        n = length(gd)
        eltys = map(typeof, first)
        if any(x -> x <: AbstractVector, eltys)
            throw(ArgumentError("mixing single values and vectors in a named tuple is not allowed"))
        end
    end
    idx = isnothing(idx_agg) ? Vector{Int}(undef, n) : idx_agg
    local initialcols
    let eltys=eltys, n=n # Workaround for julia#15276
        initialcols = ntuple(i -> Tables.allocatecolumn(eltys[i], n), _ncol(first))
    end
    targetcolnames = tuple(propertynames(first)...)
    if !extrude && first isa Union{AbstractDataFrame,
                                   NamedTuple{<:Any, <:Tuple{Vararg{AbstractVector}}}}
        outcols, finalcolnames = _combine_tables_with_first!(first, initialcols, idx, 1, 1,
                                                             f, gd, incols, targetcolnames,
                                                             firstmulticol)
    else
        outcols, finalcolnames = _combine_rows_with_first!(first, initialcols, 1, 1,
                                                           f, gd, incols, targetcolnames,
                                                           firstmulticol)
    end
    return idx, outcols, collect(Symbol, finalcolnames)
end

function fill_row!(row, outcols::NTuple{N, AbstractVector},
                   i::Integer, colstart::Integer,
                   colnames::NTuple{N, Symbol}) where N
    if _ncol(row) != N
        throw(ArgumentError("return value must have the same number of columns " *
                            "for all groups (got $N and $(length(row)))"))
    end
    @inbounds for j in colstart:length(outcols)
        col = outcols[j]
        cn = colnames[j]
        local val
        try
            val = row[cn]
        catch
            throw(ArgumentError("return value must have the same column names " *
                                "for all groups (got $colnames and $(propertynames(row)))"))
        end
        S = typeof(val)
        T = eltype(col)
        if S <: T || promote_type(S, T) <: T
            col[i] = val
        else
            return j
        end
    end
    return nothing
end

function _combine_rows_with_first!(first::Union{NamedTuple, DataFrameRow},
                                   outcols::NTuple{N, AbstractVector},
                                   rowstart::Integer, colstart::Integer,
                                   f::Base.Callable, gd::GroupedDataFrame,
                                   incols::Union{Nothing, AbstractVector, Tuple, NamedTuple},
                                   colnames::NTuple{N, Symbol},
                                   firstmulticol::Val) where N
    len = length(gd)
    gdidx = gd.idx
    starts = gd.starts
    ends = gd.ends

    # handle empty GroupedDataFrame
    len == 0 && return outcols, colnames

    # Handle first group
    j = fill_row!(first, outcols, rowstart, colstart, colnames)
    @assert j === nothing # eltype is guaranteed to match
    # Handle remaining groups
    @inbounds for i in rowstart+1:len
        row = wrap_row(do_call(f, gdidx, starts, ends, gd, incols, i), firstmulticol)
        j = fill_row!(row, outcols, i, 1, colnames)
        if j !== nothing # Need to widen column type
            local newcols
            let i = i, j = j, outcols=outcols, row=row # Workaround for julia#15276
                newcols = ntuple(length(outcols)) do k
                    S = typeof(row[k])
                    T = eltype(outcols[k])
                    U = promote_type(S, T)
                    if S <: T || U <: T
                        outcols[k]
                    else
                        copyto!(Tables.allocatecolumn(U, length(outcols[k])),
                                1, outcols[k], 1, k >= j ? i-1 : i)
                    end
                end
            end
            return _combine_rows_with_first!(row, newcols, i, j,
                                             f, gd, incols, colnames, firstmulticol)
        end
    end
    return outcols, colnames
end

# This needs to be in a separate function
# to work around a crash due to JuliaLang/julia#29430
if VERSION >= v"1.1.0-DEV.723"
    @inline function do_append!(do_it, col, vals)
        do_it && append!(col, vals)
        return do_it
    end
else
    @noinline function do_append!(do_it, col, vals)
        do_it && append!(col, vals)
        return do_it
    end
end

function append_rows!(rows, outcols::NTuple{N, AbstractVector},
                      colstart::Integer, colnames::NTuple{N, Symbol}) where N
    if !isa(rows, Union{AbstractDataFrame, NamedTuple{<:Any, <:Tuple{Vararg{AbstractVector}}}})
        throw(ArgumentError(ERROR_ROW_COUNT))
    elseif _ncol(rows) != N
        throw(ArgumentError("return value must have the same number of columns " *
                            "for all groups (got $N and $(_ncol(rows)))"))
    end
    @inbounds for j in colstart:length(outcols)
        col = outcols[j]
        cn = colnames[j]
        local vals
        try
            vals = getproperty(rows, cn)
        catch
            throw(ArgumentError("return value must have the same column names " *
                                "for all groups (got $colnames and $(propertynames(rows)))"))
        end
        S = eltype(vals)
        T = eltype(col)
        if !do_append!(S <: T || promote_type(S, T) <: T, col, vals)
            return j
        end
    end
    return nothing
end

function _combine_tables_with_first!(first::Union{AbstractDataFrame,
                                     NamedTuple{<:Any, <:Tuple{Vararg{AbstractVector}}}},
                                     outcols::NTuple{N, AbstractVector},
                                     idx::Vector{Int}, rowstart::Integer, colstart::Integer,
                                     f::Base.Callable, gd::GroupedDataFrame,
                                     incols::Union{Nothing, AbstractVector, Tuple, NamedTuple},
                                     colnames::NTuple{N, Symbol},
                                     firstmulticol::Val) where N
    len = length(gd)
    gdidx = gd.idx
    starts = gd.starts
    ends = gd.ends
    # Handle first group

    @assert _ncol(first) == N
    if !isempty(colnames) && length(gd) > 0
        j = append_rows!(first, outcols, colstart, colnames)
        @assert j === nothing # eltype is guaranteed to match
        append!(idx, Iterators.repeated(gdidx[starts[rowstart]], _nrow(first)))
    end
    # Handle remaining groups
    @inbounds for i in rowstart+1:len
        rows = wrap_table(do_call(f, gdidx, starts, ends, gd, incols, i), firstmulticol)
        _ncol(rows) == 0 && continue
        if isempty(colnames)
            newcolnames = tuple(propertynames(rows)...)
            if rows isa AbstractDataFrame
                eltys = eltype.(eachcol(rows))
            else
                eltys = map(eltype, rows)
            end
            initialcols = ntuple(i -> Tables.allocatecolumn(eltys[i], 0), _ncol(rows))
            return _combine_tables_with_first!(rows, initialcols, idx, i, 1,
                                               f, gd, incols, newcolnames, firstmulticol)
        end
        j = append_rows!(rows, outcols, 1, colnames)
        if j !== nothing # Need to widen column type
            local newcols
            let i = i, j = j, outcols=outcols, rows=rows # Workaround for julia#15276
                newcols = ntuple(length(outcols)) do k
                    S = eltype(rows isa AbstractDataFrame ? rows[!, k] : rows[k])
                    T = eltype(outcols[k])
                    U = promote_type(S, T)
                    if S <: T || U <: T
                        outcols[k]
                    else
                        copyto!(Tables.allocatecolumn(U, length(outcols[k])), outcols[k])
                    end
                end
            end
            return _combine_tables_with_first!(rows, newcols, idx, i, j,
                                               f, gd, incols, colnames, firstmulticol)
        end
        append!(idx, Iterators.repeated(gdidx[starts[i]], _nrow(rows)))
    end
    return outcols, colnames
end


text:
# This file was generated, do not modify it. # hide
LinMod = Pipeline(
    FeatureSelector(features=[:hp1]),
    LR()
);

max_issues_repo_name: alan-turing-institute/MLJTutorials

max_forks_repo_name: alan-turing-institute/MLJTutorials


text:
using ShaderToy
shadertoy("rayprimitive.frag")


text:
#= Load these if running these tests in terminal
import Pkg; Pkg.activate("..")  # CytofRepFAM
using CytofRepFAM, Random, BSON, Test, Distributions
using Distributed
include("repfamFS_tests.jl")
=#
using CytofRepFAM
using Distributed
using DelimitedFiles

printstyled("Test fitting repFAM on simulated data with PT...\n", color=:yellow)
@testset "repFAM PT" begin
  config = init_state_const_data(N=[3,2]*100, L=Dict(0=>1, 1=>1),
                                 LMCMC=Dict(0=>2, 1=>2),
                                 mus=Dict(0 => [-2.0], 1 => [2.0]),
                                 sig2=fill(0.7, 2), seed=0)
  cfs = CytofRepFAM.Model.ConstantsFS(config[:c])
  dfs = CytofRepFAM.Model.DataFS(config[:d], config[:X])
  sfs = CytofRepFAM.Model.StateFS{Float64}(config[:s], dfs)
  tfs = CytofRepFAM.Model.TunersFS(config[:t], config[:s], config[:X])

  # CytofRepFAM.Model.printConstants(cfs)

  # Fit model.
  # For algo tests
  # nmcmc = 1000
  # nburn = 100
  # maxcores = 20

  # For compile tests
  nmcmc = 3
  nburn = 3
  maxcores = 2

  if length(workers()) != maxcores
    rmprocs(filter(w -> w > 1, workers()))
    addprocs(maxcores)
    @everywhere begin
      import Pkg; Pkg.activate("../")
      using CytofRepFAM
    end
  end

  # Temperatures
  # tempers = 1.025 .^ collect(0:(maxcores-1))
  # tempers = 1.02 .^ collect(0:(maxcores-1))
  # tempers = 1.1 .^ collect(0:(maxcores-1))
  # tempers = 1.01 .^ collect(0:(maxcores-1))  # swaps a lot
  # tempers = 2.0 .^ collect(0:(maxcores-1))
  # tempers = 1.1 .^ collect(0:(maxcores-1))
  # tempers = fill(1.0, maxcores)
  # tempers = 2.0 .^ ((collect(1:maxcores) .- 1) / (maxcores - 1))
  # tempers = 10.0 .^ ((collect(1:maxcores) .- 1) / (maxcores - 1))
  # tempers = 100.0 .^ ((collect(1:maxcores) .- 1) / (maxcores - 1))
  function gentempers(maxtemp, ntemps; degree=1)
    i = collect(1:ntemps)
    powers = (i.^ degree .- 1) / (ntemps ^ degree - 1)
    return maxtemp .^ powers
  end
  tempers = gentempers(1000, maxcores, degree=4)

  @time out = CytofRepFAM.Model.fit_fs_pt!(cfs, dfs,
                                           tempers=tempers,
                                           nmcmc=3, nburn=3,
                                           Z_marg_lamgam=0.5,
                                           Z_marg_lamgam_decay_rate=10.0,
                                           Z_marg_lamgam_min=0.05,
                                           randpair=0.8,
                                           printFreq=1, seed=0,
                                           computedden=true,
                                           computeDIC=true,
                                           computeLPML=true,
                                           time_updates=false,
                                           verbose=1)

  @time out = CytofRepFAM.Model.fit_fs_pt!(cfs, dfs,
                                           # inits=[deepcopy(sfs) for _ in tempers],
                                           swap_freq=.5,
                                           tempers=tempers,
                                           nmcmc=nmcmc, nburn=nburn,
                                           Z_marg_lamgam=0.5,
                                           Z_marg_lamgam_decay_rate=100.0,
                                           Z_marg_lamgam_min=0.05,
                                           randpair=0.33,
                                           printFreq=1, seed=0,
                                           computedden=true,
                                           computeDIC=true,
                                           computeLPML=true,
                                           time_updates=false,
                                           verbose=3)
  # rmprocs(filter(w -> w > 1, workers()))

  println("Writing Output ...")

  outdir = "results/repfam-pt"
  mkpath(outdir)

  open("$(outdir)/tempers.txt", "w") do io
    writedlm(io, tempers)
  end

  for t in 1:length(out[:lls])
    open(joinpath(outdir, "ll_$(t).txt"), "w") do io
      writedlm(io, out[:lls][t], ',')
    end
  end

  open(joinpath(outdir, "swapcounts.txt"), "w") do io
    writedlm(io, out[:swapcounts], ',')
  end

  open(joinpath(outdir, "paircounts.txt"), "w") do io
    writedlm(io, out[:paircounts], ',')
  end

  BSON.bson("$(outdir)/out_fs_pt.bson", out)
  BSON.bson("$(outdir)/data_fs_pt.bson", Dict(:simdat => config[:simdat]))
end


text:
# module Duiz

function iter(e)::Union{Void,String}
    iter_trace(e, settings[:output])
end

function iter_trace(g::Base.Generator, output::Function)::Union{Void,String}
    iter_trace(g.iter, output)
end

function iter_trace(e, output::Function)::Union{Void,String}
    out = String[]
    function proc(x)
        if string == output
            push!(out, x)
        else
            println(x)
        end
    end
    padding = 11
    proc(string(lpad("e", padding), "  = $e"))
    state = start(e)
    proc(string(lpad(state, padding), " == start(e)"))
    while true
        fin = done(e, state)
        proc(string(lpad(fin, padding) , " == $done(e, $state)"))
        fin && break
        (value, next_state) = next(e, state)
        val = string == output ? value : with_color(:green, value)
        proc(string(repeat(" ", max(0, padding - length(string(value,next_state)) - 4)),
            "(", val, ", $next_state) == next(e, $state)"))
        state = next_state
    end
    if string == output
        string(join(out, "\n"), "\n")
    else
        nothing
    end
end


text:
module MultDep

using Hecke
import Base.*

function multiplicative_group_mod_units_fac_elem(A::Array{nf_elem, 1})
  k = parent(A[1])
  @assert all(i->parent(i) === k, A)
  cp = coprime_base(A)
  M = sparse_matrix(FlintZZ)
  for a = A
    T = Tuple{Int, fmpz}[]
    for i = 1:length(cp)
      p = cp[i]
      v = valuation(a, p)
      if v != 0
        push!(T, (i, fmpz(v)))
      end
#      isone(I) && break
    end
    push!(M, sparse_row(FlintZZ, T))
  end
  h, t = Hecke.hnf_kannan_bachem(M, Val{true}, truncate = true)
  return h, t, cp
end

function *(O1::NfAbsOrd, O2::NfAbsOrd)
  k = nf(O1)
  @assert k === nf(O2)
  b1 = basis(O1, k)
  n = length(b1)
  b2 = basis(O2, k)
  p = [x*y for (x,y) in Base.Iterators.ProductIterator((b1, b2))]
  d = reduce(lcm, [denominator(x) for x = p])
  M = zero_matrix(FlintZZ, n*n, n)
  z = fmpz()
  for i = 1:n*n
    a = p[i]*d
    Hecke.elem_to_mat_row!(M, i, z, a)
  end
  h = hnf(M)
  b = nf_elem[]
  for i=1:n
    push!(b, Hecke.elem_from_mat_row(k, h, i, d))
  end
  return Hecke.Order(k, b)
end

mutable struct GeIdeal
  a::NfAbsOrdIdl
  function GeIdeal(a::NfAbsOrdIdl)
    o =order(a)
    if o.ismaximal == 1
      return new(a)
    end
    #keep track of known maximality and use this to speed things up
    o = ring_of_multipliers(a)
    if o === order(a)
      return new(a)
    else
      return new(extend(a, o))
    end
  end
  function GeIdeal(a::nf_elem)
    o = Hecke.any_order(parent(a))
    d = denominator(a, o)
    return new(o(a*d)*o), new(o(d)*o)
  end
end

import Hecke.gcd, Hecke.isone, Hecke.*, Hecke.gcd_into!, Hecke.copy, Hecke.divexact,
       Hecke.isunit, Hecke.coprime_base, Hecke.valuation

function make_compatible!(a::GeIdeal, b::GeIdeal)
  o1 = order(a.a)
  o2 = order(b.a)
  if o1 === o2
    return
  end
  o = o1*o2
  a.a = extend(a.a, o)
  b.a = extend(b.a, o)
end

#isunit, divexact, gcd, isone, *, gcd_into!, copy
isone(a::GeIdeal) = isone(a.a)
isunit(a::GeIdeal) = isone(a)
copy(a::GeIdeal) = GeIdeal(a.a)

function gcd(a::GeIdeal, b::GeIdeal)
  make_compatible!(a, b)
  return GeIdeal(a.a + b.a)
end

function gcd_into!(a::GeIdeal, b::GeIdeal, c::GeIdeal)
  make_compatible!(b, c)
  a.a = b.a + c.a
  return a
end

function *(a::GeIdeal, b::GeIdeal)
  make_compatible!(a, b)
  return GeIdeal(a.a * b.a)
end

function divexact(a::GeIdeal, b::GeIdeal)
  make_compatible!(a, b)
  return GeIdeal(divexact(a.a, b.a))
end

function coprime_base(A::Array{nf_elem, 1})
  c = Array{GeIdeal, 1}()
  for a = A
    n,d = GeIdeal(a)
    isone(n) || push!(c, n)
    isone(d) || push!(c, d)
  end
  return coprime_base(c)
end

function valuation(a::nf_elem, p::GeIdeal)
  return valuation(a, p.a)
end

end


text:
# SimLynx/src/resource.jl
# Licensed under the MIT License. See LICENSE.md file in the project root for
# full license information.

# Resources
# This is a simple resource implementation that only allows unit allocations
# and deallocations.

"A (unit) allocation of the resource to a process."
struct Allocation
    time::Float64
    process::Process
end

"""
A sharable resource with a fixed number of allocatable units. The queue
maintains a list of processes waiting for the resource.
"""
mutable struct Resource
    name::String
    units::Int64
    available::Variable{Int64}
    allocated::Variable{Int64}
    queue_length::Variable{Int64}
    queue::Array{Allocation,1}
    wait::Variable{Float64}
    function Resource(units::Int64, name::String = "resource")
        new(name,
            units,
            Variable{Int64}(units),
            Variable{Int64}(0),
            Variable{Int64}(0, history=true),
            Array{Allocation,1}(undef, 0),
            Variable{Float64}(0.0, history=true))
    end
    function Resource(name::String = "resource")
        new(name,
            typemax(Int64),
            Variable{Int64}(typemax(Int64), data=:none),
            Variable{Int64}(0, history=true),
            Variable{Int64}(0, data=:none),
            Array{Allocation,1}(undef, 0),
            Variable{Float64}(0.0, data=:none))
    end
end

Base.show(io::IO, res::Resource) =
    print(io, "Resource $(res.name)")

"Allocate a unit of the resource to a process."
function allocate(resource::Resource, process::Process)
    if current_trace()
        @printf("%9.3f: %s acquired %s\n", current_time(),
                process, resource)
    end
    if resource.available.value < typemax(Int64)
        decrement!(resource.available, 1)
    end
    increment!(resource.allocated, 1)
end

"Deallocate a unit of the resource from a process."
function deallocate(resource::Resource, process::Process)
    if current_trace()
        @printf("%9.3f: %s released %s\n", current_time(),
                process, resource)
    end
    if resource.available.value < typemax(Int64) && resource.available.value < resource.units
        increment!(resource.available, 1)
    end
    decrement!(resource.allocated, 1)
end

"""
    request(resource::Resource)

Request a unit of the resource. If a unit of the resource is not available, then
queue the request.
"""
function request(resource::Resource)
    process = current_process()
    if resource.available.value > 0
        allocate(resource, process)
        resource.wait.value = 0.0
    else
        process_state!(current_simulation.notice.element, :delayed)
        push!(resource.queue, Allocation(current_time(), process))
        increment!(resource.queue_length, 1)
        yieldto(current_simulation.control_task)
    end
end

"""
    release(resource::Resource)

Release a unit of the resource. If there are process queued for the resource,
then allocate a unit of the resource to the longest waiting process. Note that
this works for unit allocations.
"""
function release(resource::Resource)
    process = current_notice().element
    deallocate(resource, process)
    if length(resource.queue) > 0
        queued = popfirst!(resource.queue)
        wait_time = current_time() - queued.time
        if wait_time > 0.0
            resource.wait.value = wait_time
        end
        decrement!(resource.queue_length, 1)
        allocate(resource, queued.process)
        @schedule now queued.process
    end
end

"""
    @with_resource resource begin
        body
    end

Wrap the body in a request / release pair for the resource.
"""
macro with_resource(resource, body)
    quote
        request($(esc(resource)))
        try
            $(esc(body))
        finally
            release($(esc(resource)))
        end
    end
end


text:
"""
    test_FM() :: Nothing

Computes structure factors for ferromagnetic correlations on different lattices and checks if the maximum resides at the Γ point.
"""
function test_FM() :: Nothing

    # init test dummy
    k = get_momenta((0.0, 1.0 * pi), (0.0, 1.0 * pi), (0.0, 1.0 * pi), (10, 10, 10))

    # initialize ferromagnetic test correlations
    @testset "FM corr" begin
        for name in ["square", "cubic", "kagome", "hyperkagome"]
            l = get_lattice(name, 6, verbose = false)
            r = get_reduced_lattice("heisenberg", [[0.0]], l, verbose = false)
            s = compute_structure_factor(Float64[exp(-norm(r.sites[i].vec)) for i in eachindex(r.sites)], k, l, r)

            @test norm(k[:, argmax(abs.(s))]) ≈ 0.0
        end
    end

    return nothing
end

"""
    test_observable() :: Nothing

Consistency checks for current implementation of observables.
"""
function test_observable() :: Nothing

    # test if ferromagnetic correlations are Fourier transformed correctly
    test_FM()

    return nothing
end


text:
if !iscxxtype(FT)
    Base.show(io::IO, x::CxxBaseRef{FT}) = show(io, x[])
end

foreach(methods(CGAL._tostring)) do r
    @eval Base.show(io::IO, x::$(r.sig.parameters[2])) = print(io, _tostring(x))
end


text:
### GradientSample is used by samplers that compute (up to) the gradient of the log-target (for ex HMC)
immutable GradientSample{N<:Number,T<:AbstractFloat,V<:AbstractVector,A<:AbstractArray} <: AbstractSample{FirstOrder}
    values::A
    loglikelihood::V
    logprior::V
    gradloglikelihood::A
    gradlogprior::A
    GradientSample(v::AbstractArray{N},ll::AbstractVector{T},lp::AbstractVector{T},gl::AbstractArray{T},gp::AbstractArray{T}) = new(v,ll,lp,gl,gp)
end

function GradientSample{N<:Number,T<:AbstractFloat}(nparas::Integer,nsamples::Integer,::Type{N},::Type{T})
    v = zeros(N,_valuestuple(nparas,nsamples))
    GradientSample{N,T,Vector,Array}(v,Vector{T}(nsamples),Vector{T}(nsamples),similar(v,T),similar(v,T))
end

### Factory function
_samples{N<:Number,T<:AbstractFloat}(::Type{Val{:gradient}},nparas::Integer,nsamples::Integer,::Type{N},::Type{T}) = GradientSample(nparas,nsamples,N,T)

@inline sampletypename(s::GradientSample) = :gradient


text:
using STABLE
using Test

@testset "STABLE.jl" begin
    @test sayhi() == "Hello from STABLE."
end


text:
@testset "Computing identifiability defect" begin

    test_cases = []
    P = fmpq_mpoly
    DType = Union{P, Generic.Frac{P}}

    R, (x, a) = PolynomialRing(QQ, ["x", "a"])
    push!(
        test_cases,
        Dict(
             "ODE" => ODE{P}(Dict{P, DType}(x => x + a), Array{P, 1}()),
            "outputs" => [x],
            "answer" => 0
        )
    )

    R, (x, a) = PolynomialRing(QQ, ["x", "a"])
    push!(
        test_cases,
        Dict(
            "ODE" => ODE{P}(Dict{P, DType}(x => x), Array{P, 1}()),
            "outputs" => [x^2],
            "answer" => 1
        )
    )

    R, (x, y, a, b) = PolynomialRing(QQ, ["x", "y", "a", "b"])
    push!(
        test_cases,
        Dict(
            "ODE" => ODE{P}(Dict{P, DType}(x => a * y, y => b * x), Array{P, 1}()),
            "outputs" => [x],
            "answer" => 1
        )
    )

    R, (x, y, a, b, c, d, u) = PolynomialRing(QQ, ["x", "y", "a", "b", "c", "d", "u"])
    push!(
        test_cases,
        Dict(
             "ODE" => ODE{P}(Dict{P, DType}(
                 x => a * x - b * x * y + u,
                 y => -c * y + d * x * y
             ), [u]),
             "outputs" => [x],
             "answer" => 1
        )
    )

    R, (M, P0, P1, P2, PN, vs, vd, K1, K2, K3, K4, Kd, KI, V1, V2, V3, V4, vm, Km, ks, k1, k2) = PolynomialRing(QQ, ["M", "P0", "P1", "P2", "PN", "vs", "vd", "K1", "K2", "K3", "K4", "Kd", "KI", "V1", "V2", "V3", "V4", "vm", "Km", "ks", "k1", "k2"])
    push!(
        test_cases,
        Dict(
             "ODE" => ODE{P}(Dict{P, DType}(
                 M => vs * KI^4 // (KI^4 + PN^4) - vm * M // (Km + M),
                 P0 => ks * M  - V1 * P0 // (K1 + P0) + V2 * P1 // (K2 + P1),
                 P1 => V1 * P0 // (K1 + P0) + V4 * P2 // (K2 + P2) - P1 * (V2 // (K2 + P1) + V3 // (K3 + P1)),
                 P2 => V3 * P1 // (K3 + P1) - P2 * (V4 // (K4 + P2) + k1 + vd // (Kd + P2)) + k2 * PN,
                 PN => k1 * P2 - k2 * PN
             ), Array{P, 1}()),
             "outputs" => [PN],
             "answer" => 1
        )
    )

    R, (x, a1, a2, a3, a4) = PolynomialRing(QQ, ["x", "a1", "a2", "a3", "a4"])
    push!(
        test_cases,
        Dict(
             "ODE" => ODE{P}(Dict{P, DType}(
                 x => x + a1 + a2 + a3 + a4
             ), Array{P, 1}()),
             "outputs" => [x^5],
             "answer" => 3
        )
    )

    for case in test_cases
        ode, outputs = case["ODE"], case["outputs"]
        @time defect = compute_defect(ode, outputs)
        @test defect == case["answer"]
    end

end

max_stars_repo_name: pogudingleb/ExperimentsBound

max_issues_repo_name: pogudingleb/ExperimentsBound

max_forks_repo_name: pogudingleb/ExperimentsBound


text:
using NumericalMethodsforEngineers
using Test

@sym begin
  ClearAll(xi, yi, Y, R, sol, ytilde)
  xi = [0, 1//2, 1]
  yi = [0, a, 1]
  Y(x_) := LagrangePolynomial(xi, yi)
  #
  # Can be formulated as ytile(x) = F(x) + C1(a) * Ψ(x)
  #
  # F(x_) := 2*x^2 - x
  # C(a_) := -4a
  # Ψ(x_) := x^2 - x
  #
  #Y(x_) := F(x) + C(a)*Ψ(x)
  #
  R(x_) = Simplify(D(Y(x), x, 2) - 3*x - 4*Y(x))
  sol = Solve(Integrate(R(x), [x, 0, 1]), a)
  ytilde(x_) = Y(x) ./ Flatten(sol)
  SetJ(r, ToString(Simplify(R(x))))
  SetJ(s, "$(sol[1][1][2])")
  SetJ(t, ToString(Simplify(Expand(ytilde(x)))));
end

println("\n\nExample 7.16: y'' = 3x + 4y, y(0)=0, y(1)=1")
println("by 1-point subdomain Weighted Residual Method")
@sym Println("\nY(x) = ", Y(x), "\n")
@sym Println("R(x) = ", R(x), "\n")
@sym Println("C1 = $(-4*sol[1][1][2])\n")
@sym Println("ytilde_1pt_subdomain(x) = ", Simplify(ytilde(x)), "\n")
println("( Example 7.16 gives: ytilde = 1/16*x*(21x - 5) )")
println()

@eval rf_1pt_subdomain(x, a) = $(Meta.parse(r))
@eval a = $(Meta.parse(s))
@eval ytilde_1pt_subdomain(x) = $(Meta.parse(t))

rf_1pt_subdomain_1(x) = rf_1pt_subdomain(x, a)
@test r == "4 - 8a - 3x + 4x*(1 - 2x + 4a*(-1 + x))"
@test t == "(1/16)*x*(-5 + 21x)"
@test (quadgk(rf_1pt_subdomain_1, 0, 1))[1] < eps()


text:
#=
Samuel J Palmer
Implementation is based on:
[1]: Akima, Hiroshi (1970). "A new method of interpolation and smooth curve fitting
    based on local procedures". Journal of the ACM. 17: 589–602.

Original module by Samuel J. Palmer, see https://github.com/sp94/CubicSplines.jl
Modified by Stefan Mathis to allow for customized extrapolation and to use a
binary search for interval identification.
=#

module CubicSplines

export CubicSpline, gradient

# As described in section 2.3 of 1
function extrapolate(xs, ys, extrap)

    # Get first / last three values
    x1, x2, x3 = xs
    y1, y2, y3 = ys

    # Extrapolate the x-coordinates of two additional datapoints with (8) from [1]
    x4 = x3-x1+x2
    x5 = x3-x1+x3

    # Quadratic extrapolation of (x4,y4), (x5,y5)
    if extrap === nothing
        y4 = y3 + 2*(x4-x3)*(y3-y2)/(x3-x2) - (x4-x3)*(y2-y1)/(x2-x1)
        y5 = y4 + 2*(x5-x4)*(y4-y3)/(x4-x3) - (x5-x4)*(y3-y2)/(x3-x2)
    else
        # Use the evalpoly macro to calculate the extrapolation values
        y4 = @evalpoly(x4-x3, y3, extrap...)
        y5 = @evalpoly(x5-x3, y3, extrap...)
    end
    return x4, x5, y4, y5
end

"""
Construct an Akima spline according to [1]. In addition to the original method,
this implementation also allows for customized extrapolation at both spline
sides. The extrapolation behaviour is defined by the coefficients of a 3rd degree
polynomial which can be given as additional keyword arguments extrapl and extrapr
to the constructor. extrapl and/or extrapr should be either "nothing" or a
set of coefficients.

Case 1: Extrapolation is set to "nothing"
================================================================================
In this case, the default quadratic extrapolation according to section 2.3 of [1]
is used to construct the spline and the spline throws an error when evaluated
outside its boundaries.

Case 2: Extrapolation is set to a vector
================================================================================
The container is given as a n-element collection of values:
    extrapl = [p1, p2, p3, ..., pn]
which are used to evaluate the following function:
    y = p0 + p1*(x-x3) + p2*(x-x3)^2 + p3*(x-x3)^3 + ... + pn*(x-x3)^n
The constant p0 is equal to y3 by definition (continuous spline). x3 and y3
are the first/last datapoint values respectively.
"""
struct CubicSpline{T <: Real}
    xs::Vector{T}
    ys::Vector{T}
    ts::Vector{T}
    ps::Matrix{T} # Coefficients of the polynom for each spline section
    logx::Bool
    extrapl::Union{Nothing,Vector{T}}
    extrapr::Union{Nothing,Vector{T}}
    function CubicSpline(xs, ys; logx=false, extrapl=nothing, extrapr=nothing)

        # Check input data
        @assert length(xs) == length(ys)
        @assert issorted(xs)
        if !issorted(xs)
            idx = sortperm(xs)
            xs = xs[idx]
            ys = ys[idx]
        end
        if logx
            xs = log.(xs)
        end
        if length(xs) < 5
            throw("At least five datapoints are necessary to construct a cubic Akima spline, however only $(length(xs)) datapoints were given.")
        end

        xb, xa, yb, ya = extrapolate(xs[3:-1:1], ys[3:-1:1], extrapl)
        xy, xz, yy, yz = extrapolate(xs[end-2:end], ys[end-2:end], extrapr)
        xs = [xa; xb; xs; xy; xz]
        ys = [ya; yb; ys; yy; yz]
        ms = zeros(eltype(ys), length(xs)-1)
        for i in 1:length(ms)
            ms[i] = (ys[i+1]-ys[i]) ./ (xs[i+1]-xs[i])
        end
        ts = zeros(eltype(ys), length(xs))
        for i in 3:length(xs)-2

            # Equals (1) in [1]
            m1, m2, m3, m4 = ms[i-2:i+1]
            # As described in [1] p.591 (parentheses block), this is an arbitrary
            # convention to guarantee uniqueness of the solution
            if m1 == m2 && m3 == m4
                ts[i] = (m2+m3)/2
            else
                numer = abs(m4-m3)*m2 + abs(m2-m1)*m3
                denom = abs(m4-m3)    + abs(m2-m1)
                ts[i] = numer ./ denom
            end
        end

        # Spline coefficients
        ps = zeros(eltype(ys), length(xs)-1, 4)
        for i in 3:length(xs)-2
            x1, x2 = xs[i:i+1]
            y1, y2 = ys[i:i+1]
            t1, t2 = ts[i:i+1]
            p0 = y1
            p1 = t1
            p2 = (3(y2-y1)/(x2-x1)-2t1-t2)/(x2-x1)
            p3 = (t1+t2-2(y2-y1)/(x2-x1))/(x2-x1)^2
            ps[i,:] = [p0, p1, p2, p3]
        end

        # Create extrapolation coefficients vector
        if extrapl !== nothing
            extrapl = vcat([ys[3]],[c for c in extrapl]) # 3rd value equals first value of the original dataset
        end
        if extrapr !== nothing
            extrapr = vcat([ys[end-2]],[c for c in extrapr]) # 3rd last value equals last value of the original dataset
        end
        return new{eltype(xs)}(xs[3:end-2], ys[3:end-2], ts[3:end-2], ps[3:end-2,:], logx, extrapl, extrapr)
    end
end
# Functor to make the CubicSpline class callable
function (obj::CubicSpline)(x)
    return obj[x]
end
# Non-allocating version for vectors
function (obj::CubicSpline)(y::AbstractVector,x::AbstractVector)
    for i in 1:length(x)
        y[i] = obj[x[i]]
    end
    return y
end

# As described in section 2.2 of [1]
function Base.getindex(spline::CubicSpline, x::Real)

    if spline.logx
        x = log(x)
    end
    if isnan(x)
        return NaN
    end

    # Value sits on upper interval border
    if x == spline.xs[end]
        idx = length(spline.xs)-1

    # Find the corresponding interval
    else

        # Find the interval index
        idx = _binary_search_interval(spline.xs, x)

        # Extrapolation to the "left"
        if idx == 0
            if spline.extrapl === nothing
                throw("x too small ($x < $(spline.xs[1]))")
            else
                return @inbounds @evalpoly(x-spline.xs[1], spline.extrapl...)
            end
        end

        # Extrapolation to the "right"
        if idx == length(spline.xs)
            if spline.extrapr === nothing
                throw("x too big ($x > $(spline.xs[end]))")
            else
                return @inbounds @evalpoly(x-spline.xs[end], spline.extrapr...)
            end
        end
    end

    # Evaluate the spline segment
    return @inbounds @evalpoly(x-spline.xs[idx], spline.ps[idx,:]...)
end

function Base.getindex(spline::CubicSpline, xs::AbstractArray{<:Real,1})
    return [spline[x] for x in xs]
end

function Base.firstindex(spline::CubicSpline)
    return spline.xs[1]
end

function Base.lastindex(spline::CubicSpline)
    return spline.xs[end]
end

"""
Evaluate the spline gradient of degree one at a given point x
"""
function gradient(spline::CubicSpline, x::Real)
    return gradient(spline, x, 1)
end

"""
Generalized gradient of degree n at a given point x
"""
function gradient(spline::CubicSpline, x::Real, n::Integer)

    # Value sits on upper interval border
    if x == spline.xs[end]
        idx = length(spline.xs)-1

    # Find the corresponding interval
    else

        # Find the interval index
        idx = _binary_search_interval(spline.xs, x)

        # Extrapolation to the "left"
        if idx == 0
            if spline.extrapl === nothing
                throw("x too small ($x < $(spline.xs[1]))")
            else
                coeff = spline.extrapl
                x_ref = spline.xs[1]
            end

        # Extrapolation to the "right"
        elseif idx == length(spline.xs)
            if spline.extrapr === nothing
                throw("x too big ($x > $(spline.xs[end]))")
            else
                coeff = spline.extrapr
                x_ref = spline.xs[end]
            end
        else
            coeff = spline.ps[idx,:]
            x_ref = spline.xs[idx]

        end
    end

    # Calculate the coefficients of degree n
    coeff = _diffpoly(coeff, n, 1)

    # If the coefficients are empty, the derivative is 0.
    if isempty(coeff)
        return zero(x)
    else
        return @evalpoly(x-x_ref, coeff...)
    end
end

"""
Binary search to find the interval which contains a given number. The interval
is given as a sorted array of real numbers. Each neighbouring pair is interpreted
as an interval. The function returns the lower index of the interval containing
the given number.

If the returned index is 0, the number is located in the interval (-Inf, array[1])
If the returned index is length(array), the number is located in the interval [array[end], Inf)
"""
@inline function _binary_search_interval(array::Vector{T}, number::Real) where T <: Real

    left = 1
    right = length(array)

    @inbounds while left <= right
        center = (left+right)÷2
        if array[center] > number
            right = center - 1
        else # array[center] <= number
            left = center + 1
        end
    end
    return left - 1
end

"""
Calculate the coefficients for the differentiated polygon
"""
function _diffpoly(coeff::AbstractArray{<:Real,1}, n::Integer=1, counter::Integer=1)

    # Discard the first coefficient, since it vanishes during the differentiation process
    coeff = coeff[2:end] .* collect(1:(length(coeff)-1))

    # If the differentiation degree is larger than the counter, differentiate the
    # differentiated coefficients again.
    if n == counter
        return coeff
    else
        return _diffpoly(coeff, n, counter+1)
    end
end

end # module


text:
using IntroML
using Documenter
using Pluto

const DOCS_PATH = @__DIR__
const SRC_PATH = joinpath(DOCS_PATH, "src")
const NB_PATH = joinpath(dirname(DOCS_PATH), "notebooks")

function build_notebook(nbpath, htmlpath)
    @info "Building notebook at $nbpath to HTML file at $htmlpath"
    s = Pluto.ServerSession()
    nb = Pluto.SessionActions.open(s, nbpath; run_async=false)
    write(htmlpath, Pluto.generate_html(nb))
    return htmlpath
end

# build Pluto notebooks
for fn in readdir(NB_PATH)
    nbpath = joinpath(NB_PATH, fn)
    htmlpath = joinpath(SRC_PATH, first(splitext(fn)) * ".html")
    build_notebook(nbpath, htmlpath)
end

DocMeta.setdocmeta!(IntroML, :DocTestSetup, :(using IntroML); recursive=true)

makedocs(;
    modules=[IntroML],
    authors="Seth Axen <seth.axen@gmail.com> and contributors",
    repo="https://github.com/mlcolab/IntroML.jl/blob/{commit}{path}#{line}",
    sitename="IntroML.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://mlcolab.github.io/IntroML.jl",
        assets=String[],
    ),
    pages=["Home" => "index.md"],
)

deploydocs(; repo="github.com/mlcolab/IntroML.jl", devbranch="main", push_preview=true)


text:
export MRIGARKDecoupledImplicit
export MRIGARKIRK21aSandu,
    MRIGARKESDIRK34aSandu,
    MRIGARKESDIRK46aSandu,
    MRIGARKESDIRK23LSA,
    MRIGARKESDIRK24LSA

"""
    MRIGARKDecoupledImplicit(f!, backward_euler_solver, fastsolver, Γs, γ̂s, Q,
                             Δt, t0)

Construct a decoupled implicit MultiRate Infinitesimal General-structure
Additive Runge--Kutta (MRI-GARK) scheme to solve

```math
    \\dot{y} = f(y, t) + g(y, t)
```

where `f` is the slow tendency function and `g` is the fast tendency function;
see Sandu (2019).

The fast tendency is integrated using the `fastsolver` and the slow tendency
using the MRI-GARK scheme. Since this is a decoupled, implicit MRI-GARK there is no implicit coupling between the fast and slow tendencies.

The `backward_euler_solver` should be of type `AbstractBackwardEulerSolver` or
`LinearBackwardEulerSolver`, and is used to perform the backward Euler solves
for `y` given the slow tendency function, namely

```math
   y = z + α f(y, t; p)
```

Currently only [`LowStorageRungeKutta2N`](@ref) schemes are supported for
`fastsolver`

The coefficients defined by `γ̂s` can be used for an embedded scheme (only the
last stage is different).

The available concrete implementations are:

  - [`MRIGARKIRK21aSandu`](@ref)
  - [`MRIGARKESDIRK34aSandu`](@ref)
  - [`MRIGARKESDIRK46aSandu`](@ref)

### References
 - [Sandu2019](@cite)
"""
mutable struct MRIGARKDecoupledImplicit{
    T,
    RT,
    AT,
    Nstages,
    NΓ,
    FS,
    Nx,
    Ny,
    Nx_Ny,
    BE,
} <: AbstractODESolver
    "time step"
    dt::RT
    "time"
    t::RT
    "elapsed time steps"
    steps::Int
    "rhs function"
    slowrhs!::Any
    "backwark Euler solver"
    besolver!::BE
    "Storage for RHS during the `MRIGARKDecoupledImplicit` update"
    Rstages::NTuple{Nstages, AT}
    "Storage for the implicit solver data vector"
    Qhat::AT
    "RK coefficient matrices for coupling coefficients"
    Γs::NTuple{NΓ, SArray{Tuple{Nx, Ny}, RT, 2, Nx_Ny}}
    "RK coefficient matrices for embedded scheme"
    γ̂s::NTuple{NΓ, SArray{NTuple{1, Ny}, RT, 1, Ny}}
    "RK coefficient vector C (time scaling)"
    Δc::SArray{NTuple{1, Nstages}, RT, 1, Nstages}
    "fast solver"
    fastsolver::FS

    function MRIGARKDecoupledImplicit(
        slowrhs!,
        backward_euler_solver,
        fastsolver,
        Γs,
        γ̂s,
        Q::AT,
        dt,
        t0,
    ) where {AT <: AbstractArray}
        NΓ = length(Γs)

        T = eltype(Q)
        RT = real(T)

        # Compute the Δc coefficients (only explicit values kept)
        Δc = sum(Γs[1], dims = 2)

        # Couple of sanity checks on the assumptions of coefficients being of
        # the decoupled implicit structure of Sandu (2019)
        @assert all(isapprox.(Δc[2:2:end], 0; atol = 2 * eps(RT)))

        Δc = Δc[1:2:(end - 1)]

        # number of slow RHS values we need to keep
        Nstages = length(Δc)

        # Couple more sanity checks on the decoupled implicit structure
        @assert Nstages == size(Γs[1], 2) - 1
        @assert Nstages == div(size(Γs[1], 1), 2)

        # Scale in the Δc to the Γ and γ̂, and convert to real type
        Γs = ntuple(k -> RT.(Γs[k]), NΓ)
        γ̂s = ntuple(k -> RT.(γ̂s[k]), NΓ)

        # Convert to real type
        Δc = RT.(Δc)

        # create storage for the stage values
        Rstages = ntuple(i -> similar(Q), Nstages)
        Qhat = similar(Q)

        FS = typeof(fastsolver)
        Nx, Ny = size(Γs[1])

        # Set up the backward Euler solver with the initial value of α
        α = dt * Γs[1][2, 2]
        besolver! =
            setup_backward_Euler_solver(backward_euler_solver, Q, α, slowrhs!)
        @assert besolver! isa AbstractBackwardEulerSolver
        BE = typeof(besolver!)

        new{T, RT, AT, Nstages, NΓ, FS, Nx, Ny, Nx * Ny, BE}(
            RT(dt),
            RT(t0),
            0,
            slowrhs!,
            besolver!,
            Rstages,
            Qhat,
            Γs,
            γ̂s,
            Δc,
            fastsolver,
        )
    end
end

function updatedt!(mrigark::MRIGARKDecoupledImplicit, dt)
    @assert Δt_is_adjustable(mrigark.besolver!)
    α = dt * mrigark.Γs[1][2, 2]
    update_backward_Euler_solver!(mrigark.besolver!, mrigark.Qhat, α)
    mrigark.dt = dt
end

function dostep!(Q, mrigark::MRIGARKDecoupledImplicit, param, time::Real)
    dt = mrigark.dt
    fast = mrigark.fastsolver

    Rs = mrigark.Rstages
    Δc = mrigark.Δc

    Nstages = length(Δc)
    groupsize = 256

    slowrhs! = mrigark.slowrhs!
    Γs = mrigark.Γs
    NΓ = length(Γs)

    ts = time
    groupsize = 256

    besolver! = mrigark.besolver!
    Qhat = mrigark.Qhat
    # Since decoupled implicit methods are being used, there is an purely
    # explicit stage followed by an implicit correction stage, hence the divide
    # by two
    for s in 1:Nstages
        # Stage dt
        dts = Δc[s] * dt
        stage_end_time = ts + dts

        # initialize the slow tendency stage value
        slowrhs!(Rs[s], Q, param, ts, increment = false)

        # # advance fast solution to time stage_end_time
        γs = ntuple(k -> ntuple(j -> Γs[k][2s - 1, j] / Δc[s], s), NΓ)
        mriparam = MRIParam(param, γs, realview.(Rs[1:s]), ts, dts)
        updatetime!(mrigark.fastsolver, ts)
        solve!(Q, mrigark.fastsolver, mriparam; timeend = stage_end_time)

        # correct with implicit slow solve
        # Qhat = Q + ∑_j Σ_k Γ_{sjk} dt Rs[j] / k
        # (Divide by k arises from the integration γ_{ij}(τ) in Sandu (2019);
        # see Equation (2.2b) and Definition 2.2
        γs = ntuple(k -> ntuple(j -> dt * Γs[k][2s, j] / k, s), NΓ)
        event = Event(array_device(Q))
        event = mri_create_Qhat!(array_device(Q), groupsize)(
            realview(Qhat),
            realview(Q),
            γs,
            mriparam.Rs;
            ndrange = length(realview(Q)),
            dependencies = (event,),
        )
        wait(array_device(Q), event)

        # Solve: Q = Qhat + α fslow(Q, stage_end_time)
        α = dt * Γs[1][2s, s + 1]
        besolver!(Q, Qhat, α, param, stage_end_time)

        # update time
        ts += dts
    end
end

# Compute: Qhat = Q + ∑_j Σ_k Γ_{sjk} dt Rs[j] / k
@kernel function mri_create_Qhat!(Qhat, Q, γs, Rs)
    i = @index(Global, Linear)
    @inbounds begin
        NΓ = length(γs)
        Ns = length(γs[1])
        qhat = Q[i]

        for s in 1:Ns
            ri = Rs[s][i]
            sc = γs[1][s]
            for k in 2:NΓ
                sc += γs[k][s]
            end
            qhat += sc * ri
        end
        Qhat[i] = qhat
    end
end

"""
    MRIGARKIRK21aSandu(f!, fastsolver, Q; dt, t0 = 0)

The 2rd order, 2 stage implicit scheme from Sandu (2019).
"""
function MRIGARKIRK21aSandu(
    slowrhs!,
    backward_euler_solver,
    fastsolver,
    Q;
    dt,
    t0 = 0,
)
    #! format: off
    Γ0 = [ 1 // 1 0 // 1
          -1 // 2 1 // 2 ]
    γ̂0 = [-1 // 2 1 // 2 ]
    #! format: on
    MRIGARKDecoupledImplicit(
        slowrhs!,
        backward_euler_solver,
        fastsolver,
        (Γ0,),
        (γ̂0,),
        Q,
        dt,
        t0,
    )
end

"""
    MRIGARKESDIRK34aSandu(f!, fastsolver, Q; dt, t0=0)

The 3rd order, 4 stage decoupled implicit scheme from Sandu (2019).
"""
function MRIGARKESDIRK34aSandu(
    slowrhs!,
    backward_euler_solver,
    fastsolver,
    Q;
    dt,
    t0 = 0,
)
    T = real(eltype(Q))
    μ = acot(2 * sqrt(T(2))) / 3
    λ = 1 - cos(μ) / sqrt(T(2)) + sqrt(T(3 // 2)) * sin(μ)
    @assert isapprox(-1 + 9λ - 18 * λ^2 + 6 * λ^3, 0, atol = 2 * eps(T))

    #! format: off
    Γ0 = [
          T(1 // 3)               0                        0           0
          -λ                      λ                        0           0
          (3-10λ) / (24λ-6)       (5-18λ) / (6-24λ)        0           0
          (-24λ^2+6λ+1) / (6-24λ) (-48λ^2+12λ+1) / (24λ-6) λ           0
          (3-16λ) / (12-48λ)      (48λ^2-21λ+2) / (12λ-3)  (3-16λ) / 4 0
          -λ                      0                        0           λ
         ]
    γ̂0 = [ 0                      0                        0           0]
    #! format: on
    MRIGARKDecoupledImplicit(
        slowrhs!,
        backward_euler_solver,
        fastsolver,
        (Γ0,),
        (γ̂0,),
        Q,
        dt,
        t0,
    )
end

"""
    MRIGARKESDIRK46aSandu(f!, fastsolver, Q; dt, t0=0)

The 4th order, 6 stage decoupled implicit scheme from Sandu (2019).
"""
function MRIGARKESDIRK46aSandu(
    slowrhs!,
    implicitsolve!,
    fastsolver,
    Q;
    dt,
    t0 = 0,
)
    T = real(eltype(Q))
    μ = acot(2 * sqrt(T(2))) / 3
    λ = 1 - cos(μ) / sqrt(T(2)) + sqrt(T(3 // 2)) * sin(μ)
    @assert isapprox(-1 + 9λ - 18 * λ^2 + 6 * λ^3, 0, atol = 2 * eps(T))

    #! format: off
    Γ0 = [
                         1 // 5                             0 // 1                             0 // 1                              0 // 1                          0 // 1           0 // 1
                        -1 // 4                             1 // 4                             0 // 1                              0 // 1                          0 // 1           0 // 1
             1771023115159 // 1929363690800    -1385150376999 // 1929363690800                 0 // 1                              0 // 1                          0 // 1           0 // 1
                    914009 // 345800                 -1000459 // 345800                        1 // 4                              0 // 1                          0 // 1           0 // 1
            18386293581909 // 36657910125200       5506531089 // 80566835440       -178423463189 // 482340922700                   0 // 1                          0 // 1           0 // 1
                  36036097 // 8299200                    4621 // 118560                -38434367 // 8299200                        1 // 4                          0 // 1           0 // 1
          -247809665162987 // 146631640500800  10604946373579 // 14663164050080   10838126175385 // 5865265620032    -24966656214317 // 36657910125200             0 // 1           0 // 1
                  38519701 // 11618880               10517363 // 9682400               -23284701 // 19364800               -10018609 // 2904720                    1 // 4           0 // 1
           -52907807977903 // 33838070884800   74846944529257 // 73315820250400  365022522318171 // 146631640500800  -20513210406809 // 109973730375600  -2918009798 // 1870301537  0 // 1
                        19 // 100                         -73 // 300                         127 // 300                          127 // 300                     -313 // 300         1 // 4
         ]

    Γ1 = [
                         0 // 1                             0 // 1                               0 // 1                              0 // 1                          0 // 1           0 // 1
                         0 // 1                             0 // 1                               0 // 1                              0 // 1                          0 // 1           0 // 1
            -1674554930619 // 964681845400      1674554930619 // 964681845400                    0 // 1                              0 // 1                          0 // 1           0 // 1
                  -1007739 // 172900                  1007739 // 172900                          0 // 1                              0 // 1                          0 // 1           0 // 1
            -8450070574289 // 18328955062600     -39429409169 // 40283417720          173621393067 // 120585230675                   0 // 1                          0 // 1           0 // 1
                -122894383 // 16598400                  14501 // 237120                  121879313 // 16598400                       0 // 1                          0 // 1           0 // 1
            32410002731287 // 15434909526400  -46499276605921 // 29326328100160    -34914135774643 // 11730531240064    45128506783177 // 18328955062600             0 // 1           0 // 1
                -128357303 // 23237760              -35433927 // 19364800                 71038479 // 38729600                 8015933 // 1452360                    0 // 1           0 // 1
           136721604296777 // 67676141769600 -349632444539303 // 146631640500800 -1292744859249609 // 293263281001600    8356250416309 // 54986865187800   17282943803 // 3740603074  0 // 1
                         3 // 25                          -29 // 300                            71 // 300                           71 // 300                     -149 // 300         0 // 1
         ]

    γ̂0 = [-1 // 4 5595 // 8804 -2445 // 8804 -4225 // 8804 2205 // 4402 -567 // 4402]
    γ̂1 = [ 0 // 1    0 // 1        0 // 1        0 // 1       0 // 1       0 // 1   ]
    #! format: on
    MRIGARKDecoupledImplicit(
        slowrhs!,
        implicitsolve!,
        fastsolver,
        (Γ0, Γ1),
        (γ̂0, γ̂1),
        Q,
        dt,
        t0,
    )
end

"""
    MRIGARKESDIRK23LSA(f!, fastsolver, Q; dt, t0 = 0, δ = 0

A 2nd order, 3 stage decoupled implicit scheme. It is based on L-Stable,
stiffly-accurate ESDIRK scheme of Bank et al (1985); see also Kennedy and
Carpenter (2016).

The free parameter `δ` can take any values for accuracy.

### References
 - [Bank1985](@cite)
 - [KennedyCarpenter2016](@cite)
"""
function MRIGARKESDIRK23LSA(
    slowrhs!,
    implicitsolve!,
    fastsolver,
    Q;
    dt,
    t0 = 0,
    δ = 0,
)
    T = real(eltype(Q))
    rt2 = sqrt(T(2))

    #! format: off
    Γ0 = [
          2 - rt2                      0                       0
          (1 - rt2) / rt2              (rt2 - 1) / rt2         0
          δ                            rt2 - 1 - δ             0
          (3 - 2rt2 * (1 + δ)) / 2rt2  (δ * 2rt2 - 1) / 2rt2   (rt2 - 1) / rt2
         ]

    γ̂0 = [0 0 0]
    #! format: on
    Δc = sum(Γ0, dims = 2)

    # Check that the explicit steps match the Δc values:
    @assert Γ0[1, 1] ≈ Δc[1]
    @assert Γ0[3, 1] + Γ0[3, 2] ≈ Δc[3]

    # Check the implicit stages have no Δc
    @assert isapprox(Γ0[2, 1] + Γ0[2, 2], 0, atol = eps(T))
    @assert isapprox(Γ0[4, 1] + Γ0[4, 2] + Γ0[4, 3], 0, atol = eps(T))

    # Check consistency with the original scheme
    @assert Γ0[1, 1] + Γ0[2, 1] ≈ 1 - 1 / rt2
    @assert Γ0[2, 2] ≈ 1 - 1 / rt2

    @assert Γ0[1, 1] + Γ0[2, 1] + Γ0[3, 1] + Γ0[4, 1] ≈ 1 / (2 * rt2)
    @assert Γ0[2, 2] + Γ0[3, 2] + Γ0[4, 2] ≈ 1 / (2 * rt2)
    @assert Γ0[4, 3] ≈ 1 - 1 / rt2

    MRIGARKDecoupledImplicit(
        slowrhs!,
        implicitsolve!,
        fastsolver,
        (Γ0,),
        (γ̂0,),
        Q,
        dt,
        t0,
    )
end

"""
    MRIGARKESDIRK24LSA(f!,
                       fastsolver,
                       Q;
                       dt,
                       t0 = 0,
                       γ = 0.2,
                       c3 = (2γ + 1) / 2,
                       a32 = 0.2,
                       α = -0.1,
                       β1 = c3 / 10,
                       β2 = c3 / 10,
                       )

A 2nd order, 4 stage decoupled implicit scheme. It is based on an L-Stable,
stiffly-accurate ESDIRK.
"""
function MRIGARKESDIRK24LSA(
    slowrhs!,
    implicitsolve!,
    fastsolver,
    Q;
    dt,
    t0 = 0,
    γ = 0.2,
    c3 = (2γ + 1) / 2,
    a32 = 0.2,
    α = -0.1,
    β1 = c3 / 10,
    β2 = c3 / 10,
)
    T = real(eltype(Q))

    # Check L-Stability constraint; bound comes from Kennedy and Carpenter
    # (2016) Table 5.
    @assert 0.1804253064293985641345831 ≤ γ < 1 // 2

    # check the stage times are increasing
    @assert 2γ < c3 < 1

    # Original RK scheme
    # Enforce L-Stability
    b3 = (2 * (1 - γ)^2 - 1) / 4 / a32

    # Enforce 2nd order accuracy
    b2 = (1 - 2γ - 2b3 * c3) / 4γ

    A = [
        0 0 0 0
        γ γ 0 0
        c3 - a32-γ a32 γ 0
        1 - b2 - b3-γ b2 b3 γ
    ]
    c = sum(A, dims = 2)
    b = A[end, :]

    # Check 2nd order accuracy
    @assert sum(b) ≈ 1
    @assert 2 * sum(A' * b) ≈ 1

    # Setup the GARK Tableau
    Δc = [c[2], 0, c[3] - c[2], 0, c[4] - c[3], 0]

    Γ0 = zeros(T, 6, 4)
    Γ0[1, 1] = Δc[1]

    Γ0[2, 1] = A[2, 1] - Γ0[1, 1]
    Γ0[2, 2] = A[2, 2]

    Γ0[3, 1] = α
    Γ0[3, 2] = Δc[3] - Γ0[3, 1]

    Γ0[4, 1] = A[3, 1] - Γ0[1, 1] - Γ0[2, 1] - Γ0[3, 1]
    Γ0[4, 2] = A[3, 2] - Γ0[1, 2] - Γ0[2, 2] - Γ0[3, 2]
    Γ0[4, 3] = A[3, 3]

    Γ0[5, 1] = β1
    Γ0[5, 2] = β2
    Γ0[5, 3] = Δc[5] - Γ0[5, 1] - Γ0[5, 2]

    Γ0[6, 1] = A[4, 1] - Γ0[1, 1] - Γ0[2, 1] - Γ0[3, 1] - Γ0[4, 1] - Γ0[5, 1]
    Γ0[6, 2] = A[4, 2] - Γ0[1, 2] - Γ0[2, 2] - Γ0[3, 2] - Γ0[4, 2] - Γ0[5, 2]
    Γ0[6, 3] = A[4, 3] - Γ0[1, 3] - Γ0[2, 3] - Γ0[3, 3] - Γ0[4, 3] - Γ0[5, 3]
    Γ0[6, 4] = A[4, 4]

    γ̂0 = [0 0 0 0]

    # Check consistency with original scheme
    @assert all(A ≈ [0 0 0 0; accumulate(+, Γ0, dims = 1)[2:2:end, :]])
    @assert all(Δc ≈ sum(Γ0, dims = 2))

    MRIGARKDecoupledImplicit(
        slowrhs!,
        implicitsolve!,
        fastsolver,
        (Γ0,),
        (γ̂0,),
        Q,
        dt,
        t0,
    )
end


text:
_sfig(x...; k...) = isdefined(ProjAssistant, :ImgTools) ?
    ImgTools.sfig(x...; k...) : error("You must 'import Plots'")
_sgif(x...; k...) = isdefined(ProjAssistant, :ImgTools) ?
    ImgTools.sgif(x...; k...) : error("You must 'import Plots'")
_lfig(x...) = isdefined(ProjAssistant, :ImgTools) ?
    ImgTools.lfig(x...) : error("You must 'import Plots'")

lfig(x...) = _lfig(x...)

function sfig(p, arg, args...;
        print_fun::Function = global_conf(:PRINT_FUN),
        mkdir = global_conf(:MK_DIR),
        verbose = global_conf(:VERBOSE),
        msg::AbstractString = "",
        kwargs...
    )

    datfile = dfname(arg, args...)
    mkdir && mkpath(dirname(datfile))
    ret = _sfig(p, datfile; kwargs...)
    verbose && print_fun("FIGURE SAVED",
        isempty(msg) ? "" : string("\n", msg),
        "\ndir: ", relpath(dirname(abspath(datfile))),
        "\nfile: ", basename(datfile),
        "\nsize: ", filesize(datfile), " bytes",
        "\n"
    )
    ret
end

function sgif(p, arg, args...;
        print_fun::Function = global_conf(:PRINT_FUN),
        mkdir = global_conf(:MK_DIR),
        verbose = global_conf(:VERBOSE),
        msg::AbstractString = "",
        kwargs...
    )
    datfile = dfname(arg, args...)
    mkdir && mkpath(dirname(datfile))
    ret = _sgif(p, datfile; kwargs...)
    verbose && print_fun("GIF SAVED",
        isempty(msg) ? "" : string("\n", msg),
        "\ndir: ", relpath(dirname(abspath(datfile))),
        "\nfile: ", basename(datfile),
        "\nsize: ", filesize(datfile), " bytes",
        "\n"
    )
    ret
end


text:
mutable struct MonteCarloDiffEqModel{absdiffeqmodel <: DiffEqBase.AbstractEnsembleProblem} <: ScalarMonteCarloProcess{Float64}
    model::absdiffeqmodel
    final_trasform::Function
    underlying::AbstractUnderlying
    function MonteCarloDiffEqModel(model::absdiffeqmodel, final_trasform::Function, underlying::AbstractUnderlying) where {absdiffeqmodel <: DiffEqBase.AbstractEnsembleProblem}
        return new{absdiffeqmodel}(model, final_trasform, underlying)
    end
    function MonteCarloDiffEqModel(model::absdiffeqmodel, underlying::AbstractUnderlying) where {absdiffeqmodel <: DiffEqBase.AbstractEnsembleProblem}
        func(x) = identity(x)
        return new{absdiffeqmodel}(model, func, underlying)
    end
end

export MonteCarloDiffeEqModel;

function simulate!(X, mcProcess::MonteCarloDiffEqModel, rfCurve::ZeroRate, mcBaseData::MonteCarloConfiguration{type1, type2, type3, type4, type5}, T::numb) where {numb <: Number, type1 <: Number, type2 <: Number, type3 <: AbstractMonteCarloMethod, type4 <: BaseMode, type5 <: Random.AbstractRNG}
    Nsim = mcBaseData.Nsim
    Nstep = mcBaseData.Nstep
    Dt = T / Nstep
    @assert T > 0.0
    diffeqmodel = mcProcess.model
    sol = solve(diffeqmodel, SOSRI(), EnsembleThreads(), trajectories = Nsim, dt = Dt, adaptive = false)
    X .= [path.u[j] for path in sol.u, j = 1:(Nstep+1)]
    f(x) = mcProcess.final_trasform(x)
    broadcast!(f, X, X)
    nothing
end


text:
using Unfolding
using CSV
using DataFrames
using Test

@testset "Unfolding.jl" begin
    # Reading data
    df_samps = CSV.read("samples.csv", DataFrame)
    df_block = CSV.read("block_model.csv", DataFrame)
    println("- Samples loaded")

    # Get coordinate points as matrix
    input_block = coords(df_block, columns=["XC","YC","ZC"])
    input_samps = coords(df_samps, columns=["X","Y","Z"])

    # Get reference surface points for unfolding
    ref_surface = getreference(input_block, axis=["X","Y"])
    println("- Reference surface extracted")

    # Get transformed coordinates of blocks and samples after unfolding
    optim = (neigh=70,)
    unf_block = unfold(input_block, ref_surface, optim=optim)
    unf_samps = unfold(input_samps, input_block, unf_block, optim=optim)
    println("- Unfolding finished")

    # Write new XT, YT and ZT columns with the transformed coordinates
    for (i,c) in enumerate([:XT,:YT,:ZT])
        df_samps[:,c] = unf_samps[i,:]
        df_block[:,c] = unf_block[i,:]
    end

    to_vtk("test_out_vtk", unf_samps, (test=rand(size(unf_samps,2)),))

    all_input = hcat(input_samps,input_block)
    all_unf   = hcat(unf_samps,unf_block)
    error     = errors(:dists, all_input, all_unf)
    good, bad = errors(:ids, all_input, all_unf)

    @test length(good) > (30 * length(bad))

    println("- Quick test of unfold with multiple inputs")
    s1, s2 = unfold([input_samps[:,1:80],input_samps[:,80:end]], ref_surface)
    s1, s2 = unfold([input_samps[:,1:80],input_samps[:,80:end]], input_block, unf_block)
end


text:
using PackageCompiler
import PackageCompiler:march, default_app_cpu_target
march()=nothing
default_app_cpu_target() = "arm1176jzf-s"
create_app(".", "build_rpizero", force=true)


text:
using Distributions: Uniform

P1 = BallInf{Float64}([0.6, 0.1, -1.3, -0.4, 0.2], 0.6)
A = [1.0  0.0;
     0.0  1.0;
    -1.0  0.0;
     0.0 -1.0]
b = [1.0, 2.0, 3.0, 4.0]
P2 = HPolyhedron(A, b)
P3 = HPolyhedron([0.0 0.1], [3.0])
P4 = Ball2{Float64}([0.2, -0.3, -1.1, 0.6, -0.7], 0.4)

## Test rand samples are contained in the set
p1 = LazySets.sample(P1)
@test p1 ∈ P1
p1_samples = LazySets.sample(P1, 100)
@test sum(p1_samples .∈ Ref(P1)) == length(p1_samples)

@test LazySets.RejectionSampler(P2).box_approx == [Uniform(-3.0,1.0), Uniform(-4.0,2.0)]


text:
@enum MSHBlockType MSHFormatBlock MSHNodesBlock MSHElementsBlock MSHUnknownBlock

function load(fs::Stream{format"MSH"}; facetype=GLTriangleFace, pointtype=Point3f)
    #GMSH MSH format (version 4)
    #http://gmsh.info/doc/texinfo/gmsh.html#MSH-file-format
    io = stream(fs)

    faces = facetype[]
    nodes = pointtype[]
    node_tags = Int[]

    while !eof(io)
        BlockType = parse_blocktype!(io)
        if BlockType == MSHNodesBlock
            parse_nodes!(io, nodes, node_tags)
        elseif BlockType == MSHElementsBlock
            parse_elements!(io, faces)
        else
            skip_block!(io)
        end
    end

    remap_faces!(faces, node_tags)

    return Mesh(nodes, faces)
end

function parse_blocktype!(io)
    header = readline(io)
    if header == "\$MeshFormat"
        return MSHFormatBlock
    elseif header == "\$Nodes"
        return MSHNodesBlock
    elseif header == "\$Elements"
        return MSHElementsBlock
    else
        return MSHUnknownBlock
    end
end

function skip_block!(io)
    while true
        line = readline(io)
        if line[1:4] == "\$End"
            break
        end
    end
    return nothing
end

function parse_nodes!(io, nodes, node_tags)
    entity_blocks, num_nodes, min_node_tag, max_node_tag = parse.(Int, split(readline(io)))
    for index_entity in 1:entity_blocks
        dim, tag, parametric, nodes_in_block = parse.(Int, split(readline(io)))
        for i in 1:nodes_in_block
            push!(node_tags, parse(eltype(node_tags), readline(io)))
        end
        for i in 1:nodes_in_block
            xyz = parse.(Float64, split(readline(io)))
            push!(nodes, xyz)
        end
    end
    endblock = readline(io)
    if endblock != "\$EndNodes"
        error("expected end block tag, got $endblock")
    end
    return nodes, node_tags
end

function parse_elements!(io, faces::Vector{T}) where T <: TriangleFace

    num_elements = parse.(Int, split(readline(io)))

    num_entity_blocks, num_elements, min_element_tag, max_element_tag = num_elements

    for index_entity in 1:num_entity_blocks

        dim, tag, element_type, elements_in_block = parse.(Int, split(readline(io)))
        if element_type == 2 # Triangles
            for i in 1:elements_in_block
                tag, n1, n2, n3 = parse.(Int, split(readline(io)))
                push!(faces, (n1, n2, n3))
            end
        else
            # for now we ignore all other elements (points, lines, hedrons, etc)
            for i in 1:elements_in_block
                readline(io)
            end
        end
    end
    endblock = readline(io)
    if endblock != "\$EndElements"
        error("expected end block tag, got $endblock")
    end
    return faces
end

function remap_faces!(faces, node_tags)
    node_map = indexin(1:maximum(node_tags), node_tags)
    for (i, face) in enumerate(faces)
        faces[i] = (node_map[face[1]], node_map[face[2]], node_map[face[3]])
    end
    return faces
end


text:
using SnpArrays
using LinearAlgebra
using Statistics
using DelimitedFiles

"""
    prune(plinkfile::AbstractString, threshold::Real, outfile::AbstractString)

Filters a PLINK file and remove all SNPs whose correlation with other SNPs
exceed `threshold`. Saves resulting plink file into `outfile`
"""
function prune(chr::Int, threshold::Real, outfile::AbstractString, ko_outfile::AbstractString)
    0 ≤ threshold ≤ 1 || error("threshold must be between 0 and 1 but was $threshold")
    plinkfile = "/scratch/users/bbchu/ukb/subset/ukb.10k.chr$chr"
    x = convert(Matrix{Float64}, SnpArray(plinkfile * ".bed"), center=true, scale=true)
    R = cor(x)
    # find all SNPs whose pairwise correlation exceeds threshold
    snps_to_exclude = Int[]
    for j in 1:size(R, 2)
        j ∈ snps_to_exclude && continue
        for i in j+1:size(R, 1)
            abs(R[i, j]) ≥ threshold && push!(snps_to_exclude, i)
        end
    end
    unique!(snps_to_exclude)
    # filter original PLINK file
    sample_idx = 1:size(x, 1) # keep all samples
    snp_idx = setdiff(1:size(x, 2), snps_to_exclude)
    SnpArrays.filter(plinkfile, sample_idx, snp_idx, des=outfile)
    writedlm("snps_kept.txt", snp_idx)
    # filter knockoff genotypes
    ko_plinkfile = "/scratch/users/bbchu/ukb/subset/ukb.10k.merged.chr$chr"
    original = vec(readdlm("/scratch/users/bbchu/ukb/subset/ukb.chr$chr.original.snp.index", Int))
    knockoff = vec(readdlm("/scratch/users/bbchu/ukb/subset/ukb.chr$chr.knockoff.snp.index", Int))
    ko_snp_idx = original[snp_idx] ∪ knockoff[snp_idx]
    SnpArrays.filter(ko_plinkfile, sample_idx, ko_snp_idx, des=ko_outfile)
    # also save which ones are knockoffs which are original SNPs
    new_original = Int[]
    new_knockoff = Int[]
    offset = 0
    for i in 1:size(x, 2)
        if i in snps_to_exclude
            offset += 2
            continue
        end
        push!(new_original, original[i] - offset)
        push!(new_knockoff, knockoff[i] - offset)
    end
    writedlm("/scratch/users/bbchu/ukb/low_LD/ukb.chr$chr.original.snp.index.lowLD", new_original)
    writedlm("/scratch/users/bbchu/ukb/low_LD/ukb.chr$chr.knockoff.snp.index.lowLD", new_knockoff)
end

#
# First filter 10k subset of UKB chr10
#
chr = 10
threshold = 0.7
outfile = "/scratch/users/bbchu/ukb/low_LD/ukb.10k.lowLD.chr$chr.threshold$threshold"
ko_outfile = "/scratch/users/bbchu/ukb/low_LD/ukb.10k.lowLD.chr$chr.threshold$threshold.knockoff"
prune(chr, threshold, outfile, ko_outfile)

#
# also need to filter test data the same way
#
snp_idx = Int.(vec(readdlm("snps_kept.txt")))
cd("/scratch/users/bbchu/ukb/populations/chr10/lowLD")

populations = ["african", "asian", "bangladeshi", "british", "caribbean", "chinese",
    "indian", "irish", "pakistani", "white_asian", "white_black", "white"]
for pop in populations
    plinkfile = "/scratch/users/bbchu/ukb/populations/chr10/ukb.chr$chr.$pop"
    outfile = "/scratch/users/bbchu/ukb/populations/chr10/lowLD/ukb.chr$chr.$pop.lowLD"
    x = SnpArray(plinkfile * ".bed")
    sample_idx = 1:size(x, 1)
    SnpArrays.filter(plinkfile, sample_idx, snp_idx, des=outfile)
end
writedlm("snps_kept.txt", snp_idx)

#
# check correlation of PCs with SNPs
#
using SnpArrays
using LinearAlgebra
using Statistics
using DelimitedFiles
x = convert(Matrix{Float64}, SnpArray("/scratch/users/bbchu/ukb/subset/ukb.10k.chr10.bed"), center=true, scale=true)
z = readdlm("/scratch/users/bbchu/ukb/subset_pca/ukb.10k.chr10.projections.txt")

# many SNPs are highly correlated with PC1, and
# these SNPs are usually clumped together
pc1 = z[:, 1]
counter = 0
for i in 1:size(x, 2)
    if abs(cor(pc1, @view(x[:, i]))) ≥ 0.75
        counter += 1
        println("snp $i")
    end
end
counter # 8 SNPs from pos 15333 to 15346

# 0 SNPs in different population are correlated with PC1
x2 = convert(Matrix{Float64}, SnpArray("/scratch/users/bbchu/ukb/populations/chr10/ukb.chr10.african.bed"), center=true, scale=true)
pc1_sub = pc1[1:size(x2, 1)]
counter = 0
for i in 1:size(x2, 2)
    if abs(cor(pc1_sub, @view(x2[:, i]))) ≥ 0.1
        counter += 1
        println("snp $i")
    end
end
counter # 0: no SNPs >0.1 correlated with PC1 subset

#
# Seed 1111: lasso knockoff predicts better than lasso, if PC1 is confounder
# Is this because LASSO had false positives associated with PC1?
# Need to confirm with other seed
#
pc1 = z[:, 1]
counter = 0
correlated_snps = Int[]
for i in 1:size(xla, 2)
    if abs(cor(pc1, @view(xla[:, i]))) ≥ 0.75
        counter += 1
        push!(correlated_snps, i)
    end
end
correlated_snps # 8 SNPs from pos 15333 to 15346

correlated_snps ∩ correct_snps

using UnicodePlots, DelimitedFiles
original = vec(readdlm("/scratch/users/bbchu/ukb/subset/ukb.chr10.original.snp.index", Int))
knockoff = vec(readdlm("/scratch/users/bbchu/ukb/subset/ukb.chr10.knockoff.snp.index", Int))
sim = 2
ko_β_lasso = vec(readdlm("/scratch/users/bbchu/ukb/prs/confound_pc/1PC_effect0.2_Radj100/fdr0.1/sim$sim/lasso.knockoff.beta"))
W = abs.(ko_β_lasso[original]) .- abs.(ko_β_lasso[knockoff])
histogram(W)

sim = 5
ko_β_lasso = vec(readdlm("/scratch/users/bbchu/ukb/prs/confound_pc/1PC_effect0.2_Radj100/fdr0.1/sim$sim/lasso.knockoff.beta"))
[ko_β_lasso[original][15333:15346] ko_β_lasso[knockoff][15333:15346]]

#
# Most real genotypes (UKB or msprime) have r2 ≈ 0.87 between original
# SNP and its knockoff. But on my own simulated genotypes, r2 is around 0.4
#
using UnicodePlots, DelimitedFiles
using SnpArrays, LinearAlgebra
using Statistics

# UKB 10k subset (single-SNP resolution)
xko_la = convert(Matrix{Float64}, SnpArray("/scratch/users/bbchu/ukb/subset/ukb.10k.merged.chr10.bed"), center=true, scale=true)
r2 = 0
for snp in 1:size(xko_la, 2) >> 1
    r2 += cor(@view(xko_la[:, 2snp]), @view(xko_la[:, 2snp - 1]))
end
r2 /= size(xko_la, 2) >> 1 # 0.8321148302625268

# low LD (filtered highly correlated snps)
xko_la = convert(Matrix{Float64}, SnpArray("/scratch/users/bbchu/ukb/low_LD/ukb.10k.lowLD.chr10.threshold0.7.knockoff.bed"), center=true, scale=true)
r2 = 0
for snp in 1:size(xko_la, 2) >> 1
    r2 += cor(@view(xko_la[:, 2snp]), @view(xko_la[:, 2snp - 1]))
end
r2 /= size(xko_la, 2) >> 1 # 0.7967258603579113

# grouped UKB data (res5)
x = SnpArray("/scratch/users/bbchu/ukb/groups/Radj5_K50_s0/ukb_gen_chr10.bed")
xko_la = convert(Matrix{Float64}, @view(x[1:10000, 1:10000]), center=true, scale=true)
r2 = 0
for snp in 1:size(xko_la, 2) >> 1
    r2 += cor(@view(xko_la[:, 2snp]), @view(xko_la[:, 2snp - 1]))
end
r2 /= size(xko_la, 2) >> 1 # 0.6125851073067299

# grouped UKB data (res2)
x = SnpArray("/oak/stanford/groups/candes/ukbiobank_tmp/knockoffs/Radj2_K50_s1/ukb_gen_chr10.bed")
xko_la = convert(Matrix{Float64}, @view(x[1:10000, 1:10000]), center=true, scale=true)
r2 = 0
for snp in 1:size(xko_la, 2) >> 1
    r2 += cor(@view(xko_la[:, 2snp]), @view(xko_la[:, 2snp - 1]))
end
r2 /= size(xko_la, 2) >> 1 # 0.37299842907393405

# my simulated (independent) genotypes
xko_la = convert(Matrix{Float64}, SnpArray("/scratch/users/bbchu/PRS_sims/no_struct/data/sim20/knockoffs/train.knockoffs_res0.bed"), center=true, scale=true)
r2 = 0
success = 0
for snp in 1:size(xko_la, 2) >> 1
    tmp = cor(@view(xko_la[:, 2snp]), @view(xko_la[:, 2snp - 1]))
    if !isnan(tmp)
        r2 += tmp
        success += 1
    end
end
r2 /= success # 0.448942588361576

# for seed 1111
idx = findall(!iszero, β)
r2 = 0.0
for i in idx
    r2 += cor(@view(xko_la[:, 2i]), @view(xko_la[:, 2i - 1]))
end
r2 /= length(idx) # 0.8457489285303439

# try letting causal snps be those that are not highly correlated with its ko
possible_causal_snp_idx = Int[]
their_correlations = Float64[]
for snp in 1:size(xla, 2)
    r2 = cor(@view(xko_la[:, 2snp]), @view(xko_la[:, 2snp - 1]))
    if r2 ≤ 0.45
        push!(possible_causal_snp_idx, snp)
        push!(their_correlations, r2)
    end
end
[possible_causal_snp_idx their_correlations]