text:
#=
	This is a file for simulation based conflict resolution time computation.
=#

using Distributions, PyPlot
push!(LOAD_PATH, "../../VICAS/src/")
push!(LOAD_PATH, "../../DeepCorrection")
push!(LOAD_PATH, "../../DeepCorrection/src")
push!(LOAD_PATH, "../TrajectoryViz")

include("../TrajectoryViz/trajViz.jl")

const BEARING_DIST =   [[0.1047 0.0489];
   						[0.3142 0.0474];
   						[0.5236 0.0474];
   						[0.7330 0.0459];
   						[0.9425 0.0425];
   						[1.1519 0.0403];
   						[1.3614 0.0365];
   						[1.5708 0.0333];
   						[1.7802 0.0304];
   						[1.9897 0.0269];
   						[2.1991 0.0237];
   						[2.4086 0.0210];
   						[2.6180 0.0194];
   						[2.8274 0.0183];
   						[3.0369 0.0181];
   						[3.2463 0.0181];
   						[3.4558 0.0183];
   						[3.6652 0.0194];
   						[3.8746 0.0210];
   						[4.0841 0.0237];
   						[4.2935 0.0269];
   						[4.5029 0.0304];
   						[4.7124 0.0333];
   						[4.9218 0.0365];
   						[5.1313 0.0403];
   						[5.3407 0.0425];
   						[5.5501 0.0459];
   						[5.7596 0.0474];
   						[5.9690 0.0474];
   						[6.1785 0.0489];]

function distance(own::Ownship, int::Ownship)
	return sqrt((own.x - int.x)^2 + (own.y - int.y)^2)
end

function checkWithinSensing(own::Ownship, int::Ownship)
	if distance(own, int) < SENSING_RANGE * 1.02
		return true
	else
		return false
	end
end

function checkDest(ac::Ownship)
	if sqrt((ac.x - ac.x_dest)^2 + (ac.y - ac.y_dest)^2) < 50
		return true
	else
		return false
	end
end

function getIntruderStartDestCoords(ownship::Ownship)
	# Use encounter model:
	intruder_dist = Categorical(vec(BEARING_DIST[:, 2]))
	intruder_angle = BEARING_DIST[:, 1][rand(intruder_dist)]
	if intruder_angle > π
		intruder_angle -= 2 * π
	end # intruder_angle ∈ [-pi, pi], centered at 0
	init_angle = norm_angle(ownship.heading + intruder_angle)
	dest_angle = norm_angle(init_angle + deg2rad(rand() * 270 + 45)) # 45 - 315 uniform?
	x_init = ownship.x + cos(init_angle) * SENSING_RANGE
	y_init = ownship.y + sin(init_angle) * SENSING_RANGE
	x_dest = ownship.x + cos(dest_angle) * SENSING_RANGE * 1.2
	y_dest = ownship.y + sin(dest_angle) * SENSING_RANGE * 1.2
	return x_init, y_init, x_dest, y_dest
end

function addIntruders!(sr::Float64, ownship::Ownship, intArr::Vector{Ownship}, CASType::Symbol, policy)
	numNewIntruders = rand(Poisson(sr))
	for i = 1:numNewIntruders
		x0, y0, x_dest, y_dest = getIntruderStartDestCoords(ownship)
		ac = initOwnship(x0, y0, x_dest, y_dest, CASType, policy)
		push!(intArr, ac)
	end
	return numNewIntruders
end

function reset(sr::Float64, CASType::Symbol, policy)
	x0 = 0.0
	y0 = 0.0
	x_dest = 10000.0
	y_dest = 0.0
	ownship = initOwnship(x0, y0, x_dest, y_dest, CASType, policy)
	intArr = Ownship[]
	addIntruders!(sr, ownship, intArr, CASType, policy)
	return ownship, intArr
end

function step!(sr::Float64, ownship::Ownship, intArr::Vector{Ownship}, CASType::Symbol, policy)
	stepResTime = length(intArr) * dt

	acArr = vcat(ownship, intArr)

	for i in 1 : length(acArr)
		intCoords = Vector{Float64}()
		for j in 1 : length(acArr)
			if j != i
				intCoords = vcat(intCoords, [acArr[j].x, acArr[j].y, acArr[j].heading])
			end
		end

		numInt = length(intArr)
		if numInt > 0 && CASType != :NOCAS
			# Get state
			numObservedInt = getState!(acArr[i], numInt, intCoords)
			if numObservedInt > 0
				# Get action and advisory
				getAdvisory!(acArr[i])
			else
				acArr[i].advisory = :COC
			end
		else
			acArr[i].advisory = :COC
		end
	end

	# Update dynamics
	for ac in acArr
		updateOwnshipCoords!(ac)
	end

	# Check within sensing and check dest for intruders
	i = 1
	while i <= length(intArr)
		if checkDest(intArr[i]) || !checkWithinSensing(ownship, intArr[i])
		   deleteat!(intArr, i)
		end
		i += 1
	end

	numNewIntruders = addIntruders!(sr, ownship, intArr, CASType, policy)

	return stepResTime, numNewIntruders
end

function render_env(sr::Float64, t::Int64, ownship::Ownship, intruders::Vector{Ownship})
	All_AC = vcat(ownship, intruders)
	AIRSPACE_DIM = 1500
	num_sections = ownship.policy.env.num_sections
	for ac in All_AC
		if ac.advisory == :COC
			scatter(ac.x, ac.y, marker="o", color="blue", s=12)
		else
			scatter(ac.x, ac.y, marker="o", color="red", s=12)
		end

		arrow_len = AIRSPACE_DIM / 6
		arrow(
			ac.x, ac.y,
			cos(ac.heading) * arrow_len * 0.8,
			sin(ac.heading) * arrow_len * 0.8,
			head_width=AIRSPACE_DIM / 20,
			width=1,
			head_length=AIRSPACE_DIM / 20,
			overhang=0.5,
			head_starts_at_zero="true",
			facecolor="black",
			length_includes_head="true")

		scatter(ac.x_dest, ac.y_dest, marker=",", color="magenta", s=12)

		PyPlot.plot([ac.x; ac.x_dest], [ac.y; ac.y_dest],
			linestyle="--", color="black", linewidth=0.5)
	end

	# alert circle
	th = linspace(-pi, pi, 30)
	if ownship.advisory == :COC
		PyPlot.plot(ownship.x + SENSING_RANGE * cos.(th),
			ownship.y + SENSING_RANGE * sin.(th),
			linestyle="--", color="green", linewidth=0.4)
	else
		PyPlot.plot(ownship.x + SENSING_RANGE * cos.(th),
			ownship.y + SENSING_RANGE * sin.(th),
			linestyle="--", color="red", linewidth=0.4)
	end

	# NMAC circle
	PyPlot.plot(ownship.x + NMAC_RANGE * cos.(th),
			ownship.y + NMAC_RANGE * sin.(th),
			linestyle="--", color="red", linewidth=0.4)

	# plot sections and highlight the closest intruder in each section
	section_angle = 2 * pi / num_sections
	for i in 0 : num_sections / 2 - 1
		th_1 = i * section_angle + ownship.heading
		th_2 = th_1 + pi
		PyPlot.plot(ownship.x + SENSING_RANGE * [cos(th_1), cos(th_2)], ownship.y + SENSING_RANGE * [sin(th_1), sin(th_2)],
			linestyle="--", color="red", linewidth=0.4)
	end

	# for i in 1 : env.num_sections
	# 	rho = env.full_state[(i - 1) * STATE_DIM + IND_rho]
	# 	if rho != TERM_VAR
	# 		theta = env.full_state[(i - 1) * STATE_DIM + IND_theta]
	# 		scatter(rho * cos(theta), rho * sin(theta),
	# 			marker="o", color="green", s=40, linewidths=0.6)
	# 	end
	# end

	xlabel("X (m)")
	ylabel("Y (m)")
	title(string(sr) * ", " * string(t) * ": " * string(length(intruders)))
	axis("equal")
	# ax = gca()
	# ax[:set_xlim]([-AIRSPACE_DIM * 1.5, AIRSPACE_DIM * 1.5])
	# ax[:set_ylim]([-AIRSPACE_DIM * 1.5, AIRSPACE_DIM * 1.5])
	pause(0.1)
	clf()
end

function resTime(;
				spawnRateSet::Union{Vector{Int64},Int64}=[25, 35, 45],
				CASType::Symbol=:correctedSector,
				policyFile::AbstractString="",
				numEvalEposide::Int64=50,
				maxEpisodeLen::Int64=2000,
				render::Bool=false)

	println("CASType = ", string(CASType))
	avgResTimeArr = Dict{Float64, Float64}()
	stdResTimeArr = Dict{Float64, Float64}()
	problem_file = "../../DeepCorrection/logs/" * policyFile * "/final_problem.jld"
	weights_file = "../../DeepCorrection/logs/" * policyFile * "/final_weights.jld"
	policy = restore(problem_file=problem_file, weights_file=weights_file)

	for sr in spawnRateSet
		sr /= 100
		epResTimeArr = Float64[]

		for ep in 1 : numEvalEposide
			# initiate ownship and intruders (reset)
			ownship, intArr = reset(sr, CASType, policy)
			totalResTime = 0.0
			totalNumIntruders = 0
			numNewIntruders = length(intArr)
			totalNumIntruders += numNewIntruders

			for t = 1 : maxEpisodeLen
				stepResTime, numNewIntruders = step!(sr, ownship, intArr, CASType, policy)
				totalResTime += stepResTime
				totalNumIntruders += numNewIntruders
				if render
					render_env(sr, t, ownship, intArr)
				end
				if checkDest(ownship) # ownship reaches destination
					break
				end
			end # for t
			push!(epResTimeArr, totalResTime / totalNumIntruders)
		end # for ep

		avgResTimeArr[sr] = mean(epResTimeArr)
		stdResTimeArr[sr] = std(epResTimeArr)
		println("sr = ", sr, ", avgResTime = ", avgResTimeArr[sr])
	end # for sr
	return avgResTimeArr, stdResTimeArr
end

max_stars_repo_name: parachutel/DeepCorrectionCAS


text:
export ThinCategory, FreeThinCategory,
  ThinSymmetricMonoidalCategory, FreeThinSymmetricMonoidalCategory,
  Preorder, Poset, FreePreorder, El, Leq, ≤, lhs, rhs, reflexive, transitive,
  Lattice, AlgebraicLattice, meet, ∧, join, ∨, top, ⊤, bottom, ⊥

import Base: join

# Thin category
###############

""" Theory of *thin categories*

Thin categories have at most one morphism between any two objects and are
isomorphic to preorders.
"""
@theory ThinCategory{Ob,Hom} <: Category{Ob,Hom} begin
  f == g ⊣ (A::Ob, B::Ob, f::Hom(A,B), g::Hom(A,B))
end

@syntax FreeThinCategory{ObExpr,HomExpr} ThinCategory begin
  compose(f::Hom, g::Hom) = associate_unit(new(f,g; strict=true), id)
end

""" Theory of *thin symmetric monoidal category*

Thin SMCs are isomorphic to commutative monoidal prosets.
"""
@theory ThinSymmetricMonoidalCategory{Ob,Hom} <: SymmetricMonoidalCategory{Ob,Hom} begin
  f == g ⊣ (A::Ob, B::Ob, f::Hom(A,B), g::Hom(A,B))
end

@syntax FreeThinSymmetricMonoidalCategory{ObExpr,HomExpr} ThinSymmetricMonoidalCategory begin
  compose(f::Hom, g::Hom) = associate_unit(new(f,g; strict=true), id)
  otimes(A::Ob, B::Ob) = associate_unit(new(A,B), munit)
  otimes(f::Hom, g::Hom) = associate(new(f,g))
end

# Preorder
##########

""" Theory of *preorders*

The generalized algebraic theory of preorders encodes inequalities ``A≤B`` as
dependent types ```Leq(A,B)`` and the axioms of reflexivity and transitivity as
term constructors.
"""
@theory Preorder{El,Leq} begin
  El::TYPE
  Leq(lhs::El, rhs::El)::TYPE
  @op (≤) := Leq

  reflexive(A::El)::(A≤A)
  transitive(f::(A≤B), g::(B≤C))::(A≤C) ⊣ (A::El, B::El, C::El)

  f == g ⊣ (A::El, B::El, f::(A≤B), g::(A≤B))
end

@syntax FreePreorder{ObExpr,HomExpr} Preorder begin
  transitive(f::Leq, g::Leq) = associate(new(f,g; strict=true))
end

# TODO: There should be an isomorphism in the category GAT between the theories
# of preorders and thin categories.
#
# @functor F(Preorder(El, Leq))::ThinCategory(Ob, Hom) begin
#   El ↦ Ob
#   Leq ↦ Hom
#   reflexive ↦ id
#   transitive ↦ compose
# end

""" Theory of *partial orders* (posets)
"""
@theory Poset{El,Leq} <: Preorder{El,Leq} begin
  A == B ⊣ (A::El, B::El, f::(A≤B), g::(B≤A))
end

# Lattice
#########

""" Theory of *lattices* as posets

A (bounded) lattice is a poset with all finite meets and joins. Viewed as a thin
category, this means that the category has all finite products and coproducts,
hence the names for the inequality constructors in the theory. Compare with
[`CartesianCategory`](@ref) and [`CocartesianCategory`](@ref).

This is one of two standard axiomatizations of a lattice, the other being
[`AlgebraicLattice`](@ref).
"""
@theory Lattice{El,Leq} <: Poset{El,Leq} begin
  @op begin
    (∧) := meet
    (⊤) := top
    (∨) := join
    (⊥) := bottom
  end

  # Meet = greatest lower bound.
  meet(A::El, B::El)::El
  proj1(A::El, B::El)::((A ∧ B) ≤ A)
  proj2(A::El, B::El)::((A ∧ B) ≤ B)
  pair(f::(C ≤ A), g::(C ≤ B))::(C ≤ (A ∧ B)) ⊣ (A::El, B::El, C::El)

  # Top = maximum.
  top()::El
  delete(A::El)::(A ≤ ⊤())

  # Join = least upper bound.
  join(A::El, B::El)::El
  coproj1(A::El, B::El)::(A ≤ (A ∨ B))
  coproj2(A::El, B::El)::(B ≤ (A ∨ B))
  copair(f::(A ≤ C), g::(B ≤ C))::((A ∨ B) ≤ C) ⊣ (A::El, B::El, C::El)

  # Bottom = minimum.
  bottom()::El
  create(A::El)::(⊥() ≤ A)
end

""" Theory of *lattices* as algebraic structures

This is one of two standard axiomatizations of a lattice, the other being
[`Lattice`](@ref). Because the partial order is not present, this theory is
merely an algebraic theory (no dependent types).

The partial order is recovered as ``A ≤ B`` iff ``A ∧ B = A`` iff ``A ∨ B = B``.
This definition could be reintroduced into a generalized algebraic theory using
an equality type `Eq(lhs::El, rhs::El)::TYPE` combined with term constructors
``meet_leq(eq::Eq(A∧B, A))::(A ≤ B)` and `join_leq(eq::Eq(A∨B, B))::(A ≤ B)`. We
do not employ that trick here because at that point it is more convenient to
just start with the poset structure, as in [`Lattice`](@ref).
"""
@theory AlgebraicLattice{El} begin
  @op begin
    (∧) := meet
    (⊤) := top
    (∨) := join
    (⊥) := bottom
  end

  # Meet/top as idempotent, commutative, associative, unital operation.
  meet(A::El, B::El)::El
  top()::El
  (A ∧ B) ∧ C == A ∧ (B ∧ C) ⊣ (A::El, B::El, C::El)
  A ∧ ⊤() == A ⊣ (A::El)
  ⊤() ∧ A == A ⊣ (A::El)
  A ∧ B == B ∧ A ⊣ (A::El, B::El)
  A ∧ A == A ⊣ (A::El)

  # Join/bottom as idempotent, commutative, associative, unital operation.
  join(A::El, B::El)::El
  bottom()::El
  (A ∨ B) ∨ C == A ∨ (B ∨ C) ⊣ (A::El, B::El, C::El)
  A ∨ ⊥() == A ⊣ (A::El)
  ⊥() ∨ A == A ⊣ (A::El)
  A ∨ B == B ∨ A ⊣ (A::El, B::El)
  A ∨ A == A ⊣ (A::El)

  # Absorption laws.
  A ∨ (A ∧ B) == A ⊣ (A::El, B::El)
  A ∧ (A ∨ B) == A ⊣ (A::El, B::El)
end

ext: jl


text:
export PermutationMatrix

import Base: *
import LinearAlgebra: transpose , inv, mul!

@doc raw"""
`PermutationMatrix(ind::Array{T})` creates a permutation matrix with ones at coordinates `(k, ind[k]) for k = 1:length(ind)`.

A permutation matrix represents a matrix containing only zeros and ones, which basically permutes the vector or matrix it is multiplied with.
These matrices A are constrained by:
```math
    A_{ij} = \{0, 1\}\\
    ∑_{i} A_{ij} = 1\\
    ∑_{j} A_{ij} = 1
```
An example is the 3-dimensional permutation matrix
```math
A = \begin{bmatrix} 1 & 0 & 0 \\ 0 & 0 & 1 \\ 1 & 0 & 0 \end{bmatrix}
```
"""
struct PermutationMatrix{ T <: Integer } <: AbstractMatrix{T}
    ind::Vector{T}
    function PermutationMatrix(ind::Vector{T}) where { T <: Integer }
        return new{T}(ind)
    end
end

@doc raw"""
`PermutationMatrix(dim::T; switch_first::Bool=true)` Generates a random permutation matrix of size (dim x dim).
The `switch_first` argument specifies whether the first index always has to be permuted.
"""
function PermutationMatrix(dim::T; switch_first::Bool=true) where { T <: Integer }
    ind = shuffle(collect(1:dim))
    if switch_first && ind[1] == 1
        tmp = ind[2]
        ind[2] = ind[1]
        ind[1] = tmp
    end
    return PermutationMatrix(ind)
end

# extensions of base functionality
Base.eltype(::PermutationMatrix{T}) where { T } = T
function Base.size(mat::PermutationMatrix)
    nr_elements = length(mat.ind)
    return (nr_elements, nr_elements)
end
Base.size(mat::PermutationMatrix, d) = d::Integer <= 2 ? length(mat.ind) : 1
Base.length(mat::PermutationMatrix)  = prod(size(mat))

function Base.getindex(mat::PermutationMatrix, i::Int, j::Int)
    if mat.ind[i] == j
        return one(eltype(mat))
    else
        return zero(eltype(mat))
    end
end

LinearAlgebra.inv(mat::PermutationMatrix)       = adjoint(mat)

# get functions
getind(mat::PermutationMatrix)      = mat.ind
getind(mat::Adjoint{T, PermutationMatrix{T}}) where { T }       = sortperm(mat.parent.ind)

# Permutation-vector multiplication
function Base.:*(P::PermutationMatrix, v::AbstractVector)
    y = similar(v)
    mul!(y, P, v)
    return y
end

function LinearAlgebra.mul!(y::AbstractVector, P::PermutationMatrix, v::AbstractVector)
    ind = getind(P)
    @inbounds @simd for k in 1:size(P,1)
        y[k] = v[ind[k]]
    end
end

function Base.:*(P::Adjoint{T, PermutationMatrix{T}}, v::AbstractVector) where { T }
    y = similar(v)
    mul!(y, P, v)
    return y
end

function LinearAlgebra.mul!(y::AbstractVector, P::Adjoint{T, PermutationMatrix{T}}, v::AbstractVector) where { T }
    ind = getind(P.parent) # explicitly take the index of the parent as not to call sortperm
    @inbounds @simd for k in 1:size(P,1)
        y[ind[k]] = v[k]
    end
end

function Base.:*(P::Transpose{T, PermutationMatrix{T}}, v::AbstractVector) where { T }
    y = similar(v)
    mul!(y, P, v)
    return y
end

function LinearAlgebra.mul!(y::AbstractVector, P::Transpose{T, PermutationMatrix{T}}, v::AbstractVector) where { T }
    ind = getind(P.parent) # explicitly take the index of the parent as not to call sortperm
    @inbounds @simd for k in 1:size(P,1)
        y[ind[k]] = v[k]
    end
end

# Permutation-matrix multiplication
function Base.:*(P::PermutationMatrix, X::AbstractMatrix)
    @assert size(X,1) == size(X,2) "Multiplication with permutation matrices is only supported for square matrices."
    Y = similar(X)
    mul!(Y, P, X)
    return Y
end

function LinearAlgebra.mul!(Y::AbstractMatrix, P::PermutationMatrix{T}, X::AbstractMatrix) where { T }
    ind = getind(P)
    @inbounds @simd for k in 1:size(P,1)
        @inbounds @simd for ki = 1:size(P,1)
            Y[k,ki] = X[ind[k],ki]
        end
    end
end

function Base.:*(P::Adjoint{T, PermutationMatrix{T}}, X::AbstractMatrix) where { T }
    @assert size(X,1) == size(X,2) "Multiplication with permutation matrices is only supported for square matrices."
    Y = similar(X)
    mul!(Y, P, X)
    return Y
end

function LinearAlgebra.mul!(Y::AbstractMatrix, P::Adjoint{T, PermutationMatrix{T}}, X::AbstractMatrix) where { T }
    ind = getind(P.parent) # explicitly take the index of the parent as not to call sortperm
    @inbounds @simd for k in 1:size(P,1)
        @inbounds @simd for ki = 1:size(P,1)
            Y[ind[k],ki] = X[k,ki]
        end
    end
end

function Base.:*(P::Transpose{T, PermutationMatrix{T}}, X::AbstractMatrix) where { T }
    @assert size(X,1) == size(X,2) "Multiplication with permutation matrices is only supported for square matrices."
    Y = similar(X)
    mul!(Y, P, X)
    return Y
end

function LinearAlgebra.mul!(Y::AbstractMatrix, P::Transpose{T, PermutationMatrix{T}}, X::AbstractMatrix) where { T }
    ind = getind(P.parent) # explicitly take the index of the parent as not to call sortperm
    @inbounds @simd for k in 1:size(P,1)
        @inbounds @simd for ki = 1:size(P,1)
            Y[ind[k],ki] = X[k,ki]
        end
    end
end

# transposedvector - matrix product
Base.:*(x::Adjoint{T, S}, P::PermutationMatrix) where { T, S <: AbstractVector } = (P'*x')'
Base.:*(x::Adjoint{T, S}, P::Adjoint{T1, PermutationMatrix{T1}}) where { T, T1, S <: AbstractVector } = (P'*x')'
Base.:*(x::Adjoint{T, S}, P::Transpose{T1, PermutationMatrix{T1}}) where { T, T1, S <: AbstractVector } = (P'*x')'
Base.:*(x::Transpose{T, S}, P::PermutationMatrix) where { T, S <: AbstractVector } = transpose(transpose(P)*transpose(x))
Base.:*(x::Transpose{T, S}, P::Adjoint{T1, PermutationMatrix{T1}}) where { T, T1, S <: AbstractVector } = transpose(transpose(P)*transpose(x))
Base.:*(x::Transpose{T, S}, P::Transpose{T1, PermutationMatrix{T1}}) where { T, T1, S <: AbstractVector } = transpose(transpose(P)*transpose(x))

# matrix-Permutation multiplication
function Base.:*(X::AbstractMatrix, P::PermutationMatrix)
    @assert size(X,1) == size(X,2) "Multiplication with permutation matrices is only supported for square matrices."
    Y = similar(X)
    mul!(Y, X, P)
    return Y
end

function LinearAlgebra.mul!(Y::AbstractMatrix, X::AbstractMatrix, P::PermutationMatrix)
    ind = getind(P)
    @inbounds @simd for k in 1:size(P,1)
        @inbounds @simd for ki = 1:size(P,1)
            Y[ki,ind[k]] = X[ki,k]
        end
    end
end

function Base.:*(X::AbstractMatrix, P::Adjoint{T, PermutationMatrix{T}}) where { T }
    @assert size(X,1) == size(X,2) "Multiplication with permutation matrices is only supported for square matrices."
    Y = similar(X)
    mul!(Y, X, P)
    return Y
end

function LinearAlgebra.mul!(Y::AbstractMatrix, X::AbstractMatrix, P::Adjoint{T, PermutationMatrix{T}}) where { T }
    ind = getind(P.parent) # explicitly take the index of the parent as not to call sortperm
    @inbounds @simd for k in 1:size(P,1)
        @inbounds @simd for ki = 1:size(P,1)
            Y[ki,k] = X[ki,ind[k]]
        end
    end
end

function Base.:*(X::AbstractMatrix, P::Transpose{T, PermutationMatrix{T}}, ) where { T }
    @assert size(X,1) == size(X,2) "Multiplication with permutation matrices is only supported for square matrices."
    Y = similar(X)
    mul!(Y, X, P)
    return Y
end

function LinearAlgebra.mul!(Y::AbstractMatrix, X::AbstractMatrix, P::Transpose{T, PermutationMatrix{T}}, ) where { T }
    ind = getind(P.parent) # explicitly take the index of the parent as not to call sortperm
    @inbounds @simd for k in 1:size(P,1)
        @inbounds @simd for ki = 1:size(P,1)
            Y[ki,k] = X[ki,ind[k]]
        end
    end
end

# multiplication of some square matrix with the permutation matrix
function PT_X_P(X::AbstractMatrix, P::PermutationMatrix)

    # allocate output
    Y = copy(X)

    # perform permutation operation
    PT_X_P!(Y, X, P)

    # return output
    return Y

end

function PT_X_P!(Y::AbstractMatrix, X::AbstractMatrix, P::PermutationMatrix)

    # fetch indices of permutation matrix
    ind = getind(P)

    # perform permutation operation
    @inbounds @simd for k1 = 1:size(P,1)
        @inbounds @simd for k2 = 1:size(P,1)
            Y[ind[k1], ind[k2]] = X[k1,k2]
        end
    end

end


text:
function AnywhereLevelTwo()
    x = rand(:x, Normal(0.0, 1.0))
    rand(:loop, AnywhereInLoop)
    return x
end

function AnywhereLevelOne()
    x = rand(:x, Normal(0.0, 1.0))
    y = rand(:y, AnywhereLevelTwo)
    return x
end

function AnywhereTopLevel()
    x = rand(:x, Normal(0.0, 1.0))
    y = rand(:y, AnywhereLevelOne)
    return x
end

function AnywhereInLoop()
    x = rand(:q, Normal(0.0, 1.0))
    for i in 1:50
        z = rand(:q => i, Normal(0.0, 1.0))
    end
end

@testset "Constrained selections" begin

    @testset "Constrained by address" begin
        sel = selection([(:x, ) => 5.0,
                         (:x, :q, :z) => 10.0])
        @test sel == sel
        @test haskey(sel, (:x, ))
        @test haskey(sel, :x)
        @test haskey(sel, (:x, :q, :z))
        @test sel[:x, :q, :z] == 10.0

        sel2 = selection([(:z, :m) => 15.0])
        merge!(sel, sel2)
        @test sel == selection([(:x, ) => 5.0,
                                (:x, :q, :z) => 10.0,
                                (:z, :m) => 15.0])
        arr = array(sel, Float64)
        @test arr == [5.0, 10.0, 15.0]
        back = selection(sel, arr)
        @test back == sel
    end

    @testset "Unconstrained by address" begin
        sel = selection([(:x, ), (:x, :q, :z)])
        @test sel == sel
        @test haskey(sel, (:x, ))
        @test haskey(sel, :x)
        @test haskey(sel, (:x, :q, :z))

        sel2 = selection([(:z, :m)])
        merge!(sel, sel2)
        @test sel == selection([(:x, ), (:x, :q, :z), (:z, :m)])
    end

    @testset "Anywhere" begin
        anyw = anywhere([(:x, ) => 5.0,
                         (:q => 21, ) => 10.0])
        ret, cl, w = generate(anyw, AnywhereTopLevel)
        @test get_ret(cl[:x]) == 5.0
        @test get_ret(cl[:y, :x]) == 5.0
        @test get_ret(cl[:y, :y, :x]) == 5.0
        @test get_ret(cl[:y, :y, :loop, :q => 21]) == 10.0
    end

    @testset "Filtering" begin
        observations = selection([(:x, ) => 5.0,
                                  (:z, :x) => 5.0,
                                  (:z, :z, :y) => 5.0])
        filtered = filter(x -> x == :y, x -> true, observations)
        @test has_top(filtered, (:z, :z, :y))
        filtered = filter(x -> x == :x, x -> true, observations)
        @test !has_top(filtered, (:z, :z, :y))
        @test has_top(filtered, :x)
        @test has_top(filtered, (:z, :x))
    end
end


text:
using Distributions
# cdf() is equivalent to p[dist] in R
# quantile() is equivalent to q[dist] in R

# No interaction effect

Unadjusted_SD = 24.7

SD_04 = sqrt(24.7^2*(1-0.4^2))
SD_045 = sqrt(24.7^2*(1-0.45^2))
SD_05 = sqrt(24.7^2*(1-0.5^2)

max_stars_repo_name: dshehtanian/julia-learn

max_issues_repo_name: dshehtanian/julia-learn

max_forks_repo_name: dshehtanian/julia-learn


text:
using MPI
using CLIMA
using CLIMA.Mesh.Topologies
using CLIMA.Mesh.Grids
using CLIMA.DGmethods
using CLIMA.DGmethods.NumericalFluxes
using CLIMA.Diagnostics
using CLIMA.MPIStateArrays
using CLIMA.LowStorageRungeKuttaMethod
using CLIMA.ODESolvers
using CLIMA.GenericCallbacks
using CLIMA.Atmos
using CLIMA.Atmos: vars_state, vars_aux
using CLIMA.VariableTemplates
using CLIMA.MoistThermodynamics
using CLIMA.PlanetParameters
using LinearAlgebra
using StaticArrays
using Logging
using Printf
using Dates
using FileIO
using Test

const ArrayType = CLIMA.array_type()

if !@isdefined integration_testing
  const integration_testing =
    parse(Bool, lowercase(get(ENV,"JULIA_CLIMA_INTEGRATION_TESTING","false")))
end

function Initialise_DYCOMS!(state::Vars, aux::Vars, (x,y,z), t)
  FT            = eltype(state)
  xvert::FT     = z
  Rd::FT        = R_d
  # These constants are those used by Stevens et al. (2005)
  qref::FT      = FT(9.0e-3)
  q_tot_sfc::FT = qref
  q_pt_sfc      = PhasePartition(q_tot_sfc)
  Rm_sfc::FT    = 461.5 #gas_constant_air(q_pt_sfc)
  T_sfc::FT     = 292.5
  P_sfc::FT     = 101780 #MSLP
  ρ_sfc::FT     = P_sfc / Rm_sfc / T_sfc
  # Specify moisture profiles
  q_liq::FT      = 0
  q_ice::FT      = 0
  zb::FT         = 600         # initial cloud bottom
  zi::FT         = 840         # initial cloud top
  ziplus::FT     = 875
  dz_cloud       = zi - zb
  q_liq_peak::FT = 0.00045     # cloud mixing ratio at z_i
  if xvert > zb && xvert <= zi
    q_liq = (xvert - zb)*q_liq_peak/dz_cloud
  end
  if xvert <= zi
    θ_liq = FT(289)
    q_tot = qref
  else
    θ_liq = FT(297.5) + (xvert - zi)^(FT(1/3))
    q_tot = FT(1.5e-3)
  end

  # Calculate PhasePartition object for vertical domain extent
  q_pt  = PhasePartition(q_tot, q_liq, q_ice)
  Rm    = gas_constant_air(q_pt)

  # Pressure
  H     = Rm_sfc * T_sfc / grav;
  p     = P_sfc * exp(-xvert/H);
  # Density, Temperature
  T = air_temperature_from_liquid_ice_pottemp_given_pressure(θ_liq, p, q_pt)
  ρ = air_density(T, p, q_pt)

  # Assign State Variables
  u1, u2 = FT(6), FT(7)
  v1, v2 = FT(-4.25), FT(-5.5)
  w = FT(10 + 0.5 * sin(2 * π * ((x/1500) + (y/1500))))
  u = (5 + 2 * sin(2 * π * ((x/1500) + (y/1500))))
  v = FT(5 + 2 * sin(2 * π * ((x/1500) + (y/1500))))
  e_kin       = FT(1/2) * (u^2 + v^2 + w^2)
  e_pot       = grav * xvert
  E           = ρ * total_energy(e_kin, e_pot, T, q_pt)
  state.ρ     = ρ
  state.ρu    = SVector(ρ*u, ρ*v, ρ*w)
  state.ρe    = E
  state.moisture.ρq_tot = ρ * q_tot
end

let
  CLIMA.init()
  mpicomm = MPI.COMM_WORLD

  ll = uppercase(get(ENV, "JULIA_LOG_LEVEL", "INFO"))
  loglevel = ll == "DEBUG" ? Logging.Debug :
    ll == "WARN"  ? Logging.Warn  :
    ll == "ERROR" ? Logging.Error : Logging.Info
  logger_stream = MPI.Comm_rank(mpicomm) == 0 ? stderr : devnull
  global_logger(ConsoleLogger(logger_stream, loglevel))

  out_dir = get(ENV, "OUT_DIR", "output")
  mkpath(out_dir)

  # Problem type
  FT = Float64
  # DG polynomial order
  N = 4
  # SGS Filter constants
  C_smag = FT(0.15)
  LHF    = FT(115)
  SHF    = FT(15)
  C_drag = FT(0.0011)
  # User defined domain parameters
  Δx, Δy, Δz = 50, 50, 20
  xmin, xmax = 0, 1500
  ymin, ymax = 0, 1500
  zmin, zmax = 0, 1500

  grid_resolution = [Δx, Δy, Δz]
  domain_size     = [xmin, xmax, ymin, ymax, zmin, zmax]
  dim             = length(grid_resolution)

  brickrange = (grid1d(xmin, xmax, elemsize=FT(grid_resolution[1])*N),
                grid1d(ymin, ymax, elemsize=FT(grid_resolution[2])*N),
                grid1d(zmin, zmax, elemsize=FT(grid_resolution[end])*N))
  zmax = brickrange[dim][end]
  zsponge = FT(1200.0)

  topl = StackedBrickTopology(mpicomm, brickrange,
                              periodicity = (true, true, false),
                              boundary=((0,0),(0,0),(1,2)))
  dt = 0.01
  timeend =  dt

  # Grid setup (topl contains brickrange information)
  grid = DiscontinuousSpectralElementGrid(topl,
                                          FloatType = FT,
                                          DeviceArray = ArrayType,
                                          polynomialorder = N,
                                         )
  # Problem constants
  # Radiation model
  κ             = FT(85)
  α_z           = FT(1)
  z_i           = FT(840)
  D_subsidence  = FT(3.75e-6)
  ρ_i           = FT(1.13)
  F_0           = FT(70)
  F_1           = FT(22)
  # Geostrophic forcing
  f_coriolis    = FT(7.62e-5)
  u_geostrophic = FT(7.0)
  v_geostrophic = FT(-5.5)
  w_ref         = FT(0)
  u_relaxation  = SVector(u_geostrophic, v_geostrophic, w_ref)

  # Model definition
  model = AtmosModel(FlatOrientation(),
                     NoReferenceState(),
                     SmagorinskyLilly{FT}(C_smag),
                     EquilMoist(),
                     DYCOMSRadiation{FT}(κ, α_z, z_i, ρ_i, D_subsidence, F_0, F_1),
                     NoSubsidence{FT}(),
                     (Gravity(),
                      RayleighSponge{FT}(zmax, zsponge, 1, u_relaxation, 2),
                      GeostrophicForcing{FT}(f_coriolis, u_geostrophic, v_geostrophic)),
                     DYCOMS_BC{FT}(C_drag, LHF, SHF),
                     Initialise_DYCOMS!)
  # Balancelaw description
  dg = DGModel(model,
               grid,
               Rusanov(),
               CentralNumericalFluxDiffusive(),
               CentralGradPenalty())
  Q = init_ode_state(dg, FT(0))
  lsrk = LSRK54CarpenterKennedy(dg, Q; dt = dt, t0 = 0)

  # Get statistics during run
  diagnostics_time_str = replace(string(now()), ":" => ".")
  cbdiagnostics = GenericCallbacks.EveryXSimulationSteps(50) do (init=false)
    sim_time_str = string(ODESolvers.gettime(lsrk))
    gather_diagnostics(mpicomm, dg, Q, diagnostics_time_str, sim_time_str,
                       out_dir)
  end

  solve!(Q, lsrk; timeend=timeend, callbacks=(cbdiagnostics,))

  # Get statistics at the end of the run
  sim_time_str = string(ODESolvers.gettime(lsrk))
  gather_diagnostics(mpicomm, dg, Q, diagnostics_time_str, sim_time_str,
                     out_dir)

  # Check results
  mpirank = MPI.Comm_rank(mpicomm)
  if mpirank == 0
    d = load(joinpath(out_dir, "diagnostics-$(diagnostics_time_str).jld2"))
    Nqk = size(d["0.0"], 1)
    Nev = size(d["0.0"], 2)
    Err = 0
    Err1 = 0
    S = zeros(Nqk * Nev)
    S1 = zeros(Nqk * Nev)
    Z = zeros(Nqk * Nev)
    for ev in 1:Nev
      for k in 1:Nqk
        dv = Diagnostics.diagnostic_vars(d["0.0"][k,ev])
        S[k+(ev-1)*Nqk] = dv.vert_eddy_u_flx
        S1[k+(ev-1)*Nqk] = dv.u
        Z[k+(ev-1)*Nqk] = dv.z
        Err += (S[k+(ev-1)*Nqk] - 0.5)^2
        Err1 += (S1[k+(ev-1)*Nqk] - 5)^2
      end
    end
    Err = sqrt(Err / (Nqk * Nev))
    @test Err <= 2e-15
    @test Err1 <= 1e-16
  end
end

max_stars_repo_name: emassoud/CLIMA

max_issues_repo_name: emassoud/CLIMA

max_forks_repo_name: emassoud/CLIMA


text:
"""
    search_ora(category::String="")
    search_ora(title::String)

Searches for apps that are under the category.
    This will provide the titles of the apps (FUTURE: as well as links to the documentation/example problems).
    If category="" (default), then it lists all of the categories instead

## Examples

The following lists all apps under the category "production process"
```julia
matches = search_ora("production process")
```
"""
function search_ora(category::String = "")::Vector{String}
    pkg_path = dirname(@__FILE__)
    if category == ""
        my_list = [titlecase(replace(x, "_" => " "))
            for x in readdir(pkg_path) if isdir(pkg_path * "/" * x)]
        return my_list # list all folders
    end
    directory = lowercase(replace(category, " " => "_"))
    try
        return [titlecase(replace(x, "_" => " ")) for x in readdir(pkg_path * "/" * directory)]
    catch e
        if isa(e, SystemError)
            throw(ArgumentError("Category $(directory[1:(end-1)]) does not exist!"))
        else
            throw(e)
        end
    end
end


text:
function addAgentCode(g,p,platform;rem=false)

    if platform == "cpu"
        code = quote
                ic1New_ = N+Threads.atomic_add!(NV,$INT(1)) + 1
                idNew_ = Threads.atomic_add!(agentIdMax,$INT(1)) + 1
                id = idNew_
            end
    else
        code = quote
            ic1New_ = N+CUDA.atomic_add!(CUDA.pointer(NV,1),$INTCUDA(1)) + 1
            idNew_ = CUDA.atomic_add!(CUDA.pointer(agentIdMax,1),$INTCUDA(1)) + 1

            id = idNew_
            end
    end

    if rem #In case of removal, add new people to the list
        push!(code.args,:(keepList_[ic1New_] = ic1New_))
    end

    args = []
    for i in g
        if i.head == :kw
            if !(i.args[1] in p.agent.declaredSymbols["Local"] || i.args[1] in p.agent.declaredSymbols["Identity"])
                error(i.args[1], "is not a Local or Identity parameter of the agent.")
            end
        else
            error(i, " is not a valid assignation of code when declaring addAgent. A Valid one should be of the form parameterOfAgent = value")
        end

        if i.args[1] in args
           error(i.args[1]," has been declared more than once.")
        end

        if i.args[1] == :id
            error("id should not be declared when calling addAgent, it is assigned automatically.")
        end
        push!(code.args,:($(i.args[1])=$(i.args[2])))
        push!(args,i.args[1])
    end

    for i in p.agent.declaredSymbols["Local"]
        if !(i in args)
            error(i, " has not been assigned in new agent of addAgent.")
        end
    end

    for i in p.agent.declaredSymbols["Identity"]
        if !(i in args)  && (i != :id)
            error(i, " has not been assigned in new agent declared by addAgent in UpdateLocal.")
        end
    end

    code = vectorize_(p.agent,code,p)

    for i in keys(p.update["Local"])
        pos = findfirst(i .== p.agent.declaredSymbols["Local"])
        posNew = p.update["Local"][i]

        push!(code.args,:(localV[ic1New_,$pos]=localVCopy[ic1New_,$posNew]))
    end

    for i in keys(p.update["Identity"])
        pos = findfirst(i .== p.agent.declaredSymbols["Identity"])
        posNew = p.update["Identity"][i]

        push!(code.args,:(identityV[ic1New_,$pos]=identityVCopy[ic1New_,$posNew]))
    end

    code = postwalk(x->@capture(x,g0_[g1_,g2_] = g3_) && g1 == :ic1_ ? :($g0[ic1New_,$g2] = $g3) : x , code)

    return code
end

"""
    function addEventAddAgent_(code::Expr,p::Program_,platform::String)

Generate the functions for division events.
"""
function addEventAddAgent_(code::Expr,p::Program_,platform::String)

    if inexpr(code,:addAgent)

        if inexpr(code,:removeAgent) #In case of removal, add new people to the list
            code = postwalk(x->@capture(x,addAgent(g__)) ? addAgentCode(g,p,platform,rem=true) : x , code)
        else
            code = postwalk(x->@capture(x,addAgent(g__)) ? addAgentCode(g,p,platform) : x , code)
        end

        if platform == "cpu"

            push!(p.declareVar.args,:(NV = Threads.Atomic{$INT}(0)))
            push!(p.declareVar.args,:(agentIdMax = Threads.Atomic{$INT}(N)))

            push!(p.execInloop.args,:(N += NV[]; NV[] = 0))

        elseif platform == "gpu"

            push!(p.declareVar.args,:(NV = CUDA.zeros($INTCUDA,1)))
            push!(p.declareVar.args,:(agentIdMax = N .*CUDA.ones($INTCUDA,1)))

            push!(p.execInloop.args,:(N += Core.Array(NV)[1]; NV .= 0))

        end

        push!(p.args,:(NV),:(agentIdMax))

    end

    return code
end


text:
function spacecase(cur_string::AbstractString)
  _defaultcase(cur_string)
end


text:
include("../../src/preprocessing/fill_missing.jl")
using Test, DataFrames, Statistics

test_data = DataFrame(
    w = ["one", "two", "three", "four"],
    x = [1, 2, 3, missing],
    y = [15, missing, 90, 15],
    z = [-1, -20, 20, -20]
)

@testset "fill_missing" begin
    # mean, all columns
    @test isequal(fill_missing_aggregate(test_data, method=mean), DataFrame(
        w = ["one", "two", "three", "four"],
        x = [1.0, 2.0, 3.0, 2.0],
        y = [15.0, 40.0, 90.0, 15.0],
        z = [-1.0, -20.0, 20.0, -20.0]
    ))
    # mode, all columns
    @test isequal(fill_missing_aggregate(test_data, method=median), DataFrame(
        w = ["one", "two", "three", "four"],
        x = [1.0, 2.0, 3.0, 2.0],
        y = [15.0, 15.0, 90.0, 15.0],
        z = [-1.0, -20.0, 20.0, -20.0]
    ))
    # mean, subset of columns
    @test isequal(fill_missing_aggregate(test_data, [:y, :z],method=mean), DataFrame(
        w = ["one", "two", "three", "four"],
        x = [1, 2, 3, missing],
        y = [15.0, 40.0, 90.0, 15.0],
        z = [-1.0, -20.0, 20.0, -20.0]
    ))

    # scalar, correct type
    @test isequal(fill_missing_scalar(test_data, [:x, :y], scalar=100), DataFrame(
        w = ["one", "two", "three", "four"],
        x = [1, 2, 3, 100],
        y = [15, 100, 90, 15],
        z = [-1, -20, 20, -20]
    ))
    # scalar, incorrect type
    @test try
            fill_missing_scalar(test_data, [:x, :y], scalar="string")
        catch ex
            ex
        end == ArgumentError("Scalar must match variable type of ALL columns.")
end


text:
# Autogenerated wrapper script for boost_jll for aarch64-apple-darwin
export libboost_atomic, libboost_chrono, libboost_container, libboost_context, libboost_contract, libboost_coroutine, libboost_date_time, libboost_filesystem, libboost_graph, libboost_iostreams, libboost_log, libboost_log_setup, libboost_prg_exec_monitor, libboost_program_options, libboost_random, libboost_regex, libboost_serialization, libboost_stacktrace_basic, libboost_stacktrace_noop, libboost_system, libboost_thread, libboost_timer, libboost_type_erasure, libboost_unit_test_framework, libboost_wave, libboost_wserialization

JLLWrappers.@generate_wrapper_header("boost")
JLLWrappers.@declare_library_product(libboost_atomic, "@rpath/libboost_atomic.dylib")
JLLWrappers.@declare_library_product(libboost_chrono, "@rpath/libboost_chrono.dylib")
JLLWrappers.@declare_library_product(libboost_container, "@rpath/libboost_container.dylib")
JLLWrappers.@declare_library_product(libboost_context, "@rpath/libboost_context.dylib")
JLLWrappers.@declare_library_product(libboost_contract, "@rpath/libboost_contract.dylib")
JLLWrappers.@declare_library_product(libboost_coroutine, "@rpath/libboost_coroutine.dylib")
JLLWrappers.@declare_library_product(libboost_date_time, "@rpath/libboost_date_time.dylib")
JLLWrappers.@declare_library_product(libboost_filesystem, "@rpath/libboost_filesystem.dylib")
JLLWrappers.@declare_library_product(libboost_graph, "@rpath/libboost_graph.dylib")
JLLWrappers.@declare_library_product(libboost_iostreams, "@rpath/libboost_iostreams.dylib")
JLLWrappers.@declare_library_product(libboost_log, "@rpath/libboost_log.dylib")
JLLWrappers.@declare_library_product(libboost_log_setup, "@rpath/libboost_log_setup.dylib")
JLLWrappers.@declare_library_product(libboost_prg_exec_monitor, "@rpath/libboost_prg_exec_monitor.dylib")
JLLWrappers.@declare_library_product(libboost_program_options, "@rpath/libboost_program_options.dylib")
JLLWrappers.@declare_library_product(libboost_random, "@rpath/libboost_random.dylib")
JLLWrappers.@declare_library_product(libboost_regex, "@rpath/libboost_regex.dylib")
JLLWrappers.@declare_library_product(libboost_serialization, "@rpath/libboost_serialization.dylib")
JLLWrappers.@declare_library_product(libboost_stacktrace_basic, "@rpath/libboost_stacktrace_basic.dylib")
JLLWrappers.@declare_library_product(libboost_stacktrace_noop, "@rpath/libboost_stacktrace_noop.dylib")
JLLWrappers.@declare_library_product(libboost_system, "@rpath/libboost_system.dylib")
JLLWrappers.@declare_library_product(libboost_thread, "@rpath/libboost_thread.dylib")
JLLWrappers.@declare_library_product(libboost_timer, "@rpath/libboost_timer.dylib")
JLLWrappers.@declare_library_product(libboost_type_erasure, "@rpath/libboost_type_erasure.dylib")
JLLWrappers.@declare_library_product(libboost_unit_test_framework, "@rpath/libboost_unit_test_framework.dylib")
JLLWrappers.@declare_library_product(libboost_wave, "@rpath/libboost_wave.dylib")
JLLWrappers.@declare_library_product(libboost_wserialization, "@rpath/libboost_wserialization.dylib")
function __init__()
    JLLWrappers.@generate_init_header()
    JLLWrappers.@init_library_product(
        libboost_atomic,
        "lib/libboost_atomic.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_chrono,
        "lib/libboost_chrono.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_container,
        "lib/libboost_container.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_context,
        "lib/libboost_context.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_contract,
        "lib/libboost_contract.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_coroutine,
        "lib/libboost_coroutine.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_date_time,
        "lib/libboost_date_time.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_filesystem,
        "lib/libboost_filesystem.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_graph,
        "lib/libboost_graph.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_iostreams,
        "lib/libboost_iostreams.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_log,
        "lib/libboost_log.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_log_setup,
        "lib/libboost_log_setup.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_prg_exec_monitor,
        "lib/libboost_prg_exec_monitor.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_program_options,
        "lib/libboost_program_options.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_random,
        "lib/libboost_random.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_regex,
        "lib/libboost_regex.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_serialization,
        "lib/libboost_serialization.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_stacktrace_basic,
        "lib/libboost_stacktrace_basic.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_stacktrace_noop,
        "lib/libboost_stacktrace_noop.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_system,
        "lib/libboost_system.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_thread,
        "lib/libboost_thread.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_timer,
        "lib/libboost_timer.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_type_erasure,
        "lib/libboost_type_erasure.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_unit_test_framework,
        "lib/libboost_unit_test_framework.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_wave,
        "lib/libboost_wave.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libboost_wserialization,
        "lib/libboost_wserialization.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
# Arithmetic
#XYZ and LMS are linear vector spaces
typealias Linear3 Union{XYZ, LMS}
+{C<:Linear3}(a::C, b::C) = C(comp1(a)+comp1(b), comp2(a)+comp2(b), comp3(a)+comp3(b))
-{C<:Linear3}(a::C, b::C) = C(comp1(a)-comp1(b), comp2(a)-comp2(b), comp3(a)-comp3(b))
-(a::Linear3) = typeof(a)(-comp1(a), -comp2(a), -comp3(a))
*(c::Number, a::Linear3) = base_color_type(a)(c*comp1(a), c*comp2(a), c*comp3(a))
/(a::Linear3, c::Number) = base_color_type(a)(comp1(a)/c, comp2(a)/c, comp3(a)/c)

# Algorithms relating to color processing and generation

# Chromatic Adaptation / Whitebalancing
# -------------------------------------

"""
    whitebalance(c, src_white, ref_white)

Whitebalance a color.

Input a source (adopted) and destination (reference) white. E.g., if you have
a photo taken under florencent lighting that you then want to appear correct
under regular sunlight, you might do something like
`whitebalance(c, WP_F2, WP_D65)`.

Args:

- `c`: An observed color.
- `src_white`: Adopted or source white corresponding to `c`
- `ref_white`: Reference or destination white.

Returns:
  A whitebalanced color.
"""
function whitebalance{T <: Color}(c::T, src_white::Color, ref_white::Color)
    c_lms = convert(LMS, c)
    src_wp = convert(LMS, src_white)
    dest_wp = convert(LMS, ref_white)

    # This is sort of simplistic, it sets the degree of adaptation term in
    # CAT02 to 0.
    # Setting the degree of adaptation to 0 is rather odd. Wouldn’t setting
    # it to 1.0 make more sense as a temporary default value?
    ans = LMS(c_lms.l * dest_wp.l / src_wp.l,
              c_lms.m * dest_wp.m / src_wp.m,
              c_lms.s * dest_wp.s / src_wp.s)

    convert(T, ans)
end

# Simulation of color deficiency (color "blindness")
# ----------------------------

# This method is due to:
# Brettel, H., Viénot, F., & Mollon, J. D. (1997).  Computerized simulation of
# color appearance for dichromats. Josa A, 14(10), 2647–2655.
#
# These functions add to Brettel's method a parameter p in [0, 1] giving the
# degree of photopigment loss. At p = 1, the complete loss of a particular
# photopigment is simulated, where 0 < p < 1 represents partial loss.

# This is supposed to be "the brightest possible metamer of an equal-energy
# stimulus". I'm punting a bit and just calling that RGB white.
const default_brettel_neutral = convert(LMS, RGB(1, 1, 1))

# Helper function for Brettel conversions.
function brettel_abc(neutral::LMS, anchor::LMS)
    a = neutral.m * anchor.s - neutral.s * anchor.m
    b = neutral.s * anchor.l - neutral.l * anchor.s
    c = neutral.l * anchor.m - neutral.m * anchor.l
    (a, b, c)
end

"""
    protanopic(c)
    protanopic(c, p)

Convert a color to simulate protanopic color deficiency (lack of the
long-wavelength photopigment).  `c` is the input color; the optional
argument `p` is the fraction of photopigment loss, in the range 0 (no
loss) to 1 (complete loss).
"""
function protanopic{T <: Color}(q::T, p, neutral::LMS)
    q = convert(LMS, q)
    anchor_wavelen = q.s / q.m < neutral.s / neutral.m ? 575 : 475
    anchor = colormatch(anchor_wavelen)
    anchor = convert(LMS, anchor)
    (a, b, c) = brettel_abc(neutral, anchor)

    q = LMS((one(p) - p) * q.l + p * (-(b*q.m + c*q.s)/a),
            q.m,
            q.s)
    convert(T, q)
end

"""
    deuteranopic(c)
    deuteranopic(c, p)

Convert a color to simulate deuteranopic color deficiency (lack of the
middle-wavelength photopigment).  See the description of `protanopic` for detail about the arguments.
"""
function deuteranopic{T <: Color}(q::T, p, neutral::LMS)
    q = convert(LMS, q)
    anchor_wavelen = q.s / q.l < neutral.s / neutral.l ? 575 : 475
    anchor = colormatch(anchor_wavelen)
    anchor = convert(LMS, anchor)
    (a, b, c) = brettel_abc(neutral, anchor)

    q = LMS(q.l,
            (one(p) - p) * q.m + p * (-(a*q.l + c*q.s)/b),
            q.s)
    convert(T, q)
end

"""
    tritanopic(c)
    tritanopic(c, p)

Convert a color to simulate tritanopic color deficiency (lack of the
short-wavelength photogiment).  See `protanopic` for more detail about
the arguments.
"""
function tritanopic{T <: Color}(q::T, p, neutral::LMS)
    q = convert(LMS, q)
    anchor_wavelen = q.m / q.l < neutral.m / neutral.l ? 660 : 485
    anchor = colormatch(anchor_wavelen)
    anchor = convert(LMS, anchor)
    (a, b, c) = brettel_abc(neutral, anchor)

    q = LMS(q.l,
            q.m,
            (one(p) - p) * q.l + p * (-(a*q.l + b*q.m)/c))
    convert(T, q)
end

protanopic(c::Color, p)   = protanopic(c, p, default_brettel_neutral)
deuteranopic(c::Color, p) = deuteranopic(c, p, default_brettel_neutral)
tritanopic(c::Color, p)   = tritanopic(c, p, default_brettel_neutral)

protanopic(c::Color)   = protanopic(c, 1.0)
deuteranopic(c::Color) = deuteranopic(c, 1.0)
tritanopic(c::Color)   = tritanopic(c, 1.0)

Compat.@dep_vectorize_1arg Color protanopic
Compat.@dep_vectorize_1arg Color deuteranopic
Compat.@dep_vectorize_1arg Color tritanopic

# MSC - Most Saturated Colorant for given hue h
# ---------------------

"""
    MSC(h)
    MSC(h, l)

Calculates the most saturated color for any given hue `h` by
finding the corresponding corner in LCHuv space. Optionally,
the lightness `l` may also be specified.
"""
function MSC(h)

    #Corners of RGB cube
    const h0 = 12.173988685914473 #convert(LCHuv,RGB(1,0,0)).h
    const h1 = 85.872748860776770 #convert(LCHuv,RGB(1,1,0)).h
    const h2 = 127.72355046632740 #convert(LCHuv,RGB(0,1,0)).h
    const h3 = 192.17397321802082 #convert(LCHuv,RGB(0,1,1)).h
    const h4 = 265.87273498040290 #convert(LCHuv,RGB(0,0,1)).h
    const h5 = 307.72354567594960 #convert(LCHuv,RGB(1,0,1)).h

    #Wrap h to [0, 360] range]
    h = mod(h, 360)

    p=0 #variable
    o=0 #min
    t=0 #max

    #Selecting edge of RGB cube; R=1 G=2 B=3
    if h0 <= h < h1
        p=2; o=3; t=1
    elseif h1 <= h < h2
        p=1; o=3; t=2
    elseif h2 <= h < h3
        p=3; o=1; t=2
    elseif h3 <= h < h4
        p=2; o=1; t=3
    elseif h4 <= h < h5
        p=1; o=2; t=3
    elseif h5 <= h || h < h0
        p=3; o=2; t=1
    end

    col=zeros(3)

    #check if we are directly on the edge of the RGB cube (within some tolerance)
    for edge in [h0, h1, h2, h3, h4, h5]
        if edge - 200eps() < h < edge + 200eps()
            col[p] = edge in [h0, h2, h4] ? 0.0 : 1.0
            col[o] = 0.0
            col[t] = 1.0
            return convert(LCHuv, RGB(col[1],col[2],col[3]))
        end
    end

    alpha=-sind(h)
    beta=cosd(h)

    #un &vn are calculated based on reference white (D65)
    #X=0.95047; Y=1.0000; Z=1.08883
    const un=0.19783982482140777 #4.0X/(X+15.0Y+3.0Z)
    const vn=0.46833630293240970 #9.0Y/(X+15.0Y+3.0Z)

    #sRGB matrix
    const M=[0.4124564  0.3575761  0.1804375;
             0.2126729  0.7151522  0.0721750;
             0.0193339  0.1191920  0.9503041]'
    const g=2.4

    m_tx=M[t,1]
    m_ty=M[t,2]
    m_tz=M[t,3]
    m_px=M[p,1]
    m_py=M[p,2]
    m_pz=M[p,3]

    f1=(4alpha*m_px+9beta*m_py)
    a1=(4alpha*m_tx+9beta*m_ty)
    f2=(m_px+15m_py+3m_pz)
    a2=(m_tx+15m_ty+3m_tz)

    cp=((alpha*un+beta*vn)*a2-a1)/(f1-(alpha*un+beta*vn)*f2)

    #gamma inversion
    cp = cp <= 0.003 ? 12.92cp : 1.055cp^(1.0/g)-0.05
    #cp = 1.055cp^(1.0/g)-0.05

    col[p]=clamp01(cp)
    col[o]=0.0
    col[t]=1.0

    return convert(LCHuv, RGB(col[1],col[2],col[3]))
end

# Maximum saturation for given lightness and hue
# ----------------------

# Maximally saturated color for a specific hue and lightness
# is found by looking for the edge of LCHuv space.
function MSC(h,l)
    pmid=MSC(h)

    if l <= pmid.l
        pend=LCHuv(0,0,0)
    elseif l > pmid.l
        pend=LCHuv(100,0,0)
    end

    a=(pend.l-l)/(pend.l-pmid.l)
    a*(pmid.c-pend.c)+pend.c
end


text:
xml_evfile1 = path*"/SampleFiles/XML/fdsnws-event_2017-01-12T03-18-55Z.xml"
xml_evfile2 = path*"/SampleFiles/XML/ISC_2011-tohoku-oki.xml"
xml_stfile = path*"/SampleFiles/XML/fdsnws-station_2017-01-12T03-17-42Z.xml"

printstyled("  QuakeML\n", color=:light_green)

id_err = "error in Station ID creation!"
unit_err = "units don't match instrument code!"
true_id = String["3337497", "3279407", "2844986", "2559759", "2092067", "1916079", "2413"]
true_ot = DateTime("2011-03-11T05:46:23.200")
true_loc = Float64[2.2376 38.2963; 93.0144 142.498; 26.3 19.7]
true_mag = Float32[8.6, 9.1, 8.8, 8.5, 8.6, 9.0, 8.5]
true_msc = String["MW", "MW", "MW", "MW", "MW", "MW", ""]
r1 = PZResp(Complex{Float32}.([ 0.0+0.0im -981.0+1009.0im
                                0.0+0.0im -981.0-1009.0im
                                0.0+0.0im -3290.0+1263.0im
                                0.0+0.0im -3290.0-1263.0im]), rev=true)
r2 = PZResp(Complex{Float32}.([   0.0+0.0im       -0.037-0.037im
                                  0.0+0.0im       -0.037+0.037im
                                  -15.15+0.0im    -15.64+0.0im
                                  -176.6+0.0im    -97.34-400.7im
                                  -463.1-430.5im  -97.34+400.7im
                                  -463.1+430.5im  -374.8+0.0im
                                  0.0+0.0im       -520.3+0.0im
                                  0.0+0.0im       -10530.0-10050.0im
                                  0.0+0.0im       -10530.0+10050.0im
                                  0.0+0.0im       -13300.0+0.0im
                                  0.0+0.0im       -255.097+0.0im ]),rev=true)

printstyled("    file read 1\n", color=:light_green)
(EC,RR) = read_qml(xml_evfile1)
Nev = length(EC)
@test Nev == length(true_id)
for i = 1:Nev
  @test EC[i].id == true_id[i]
  @test EC[i].mag.val == true_mag[i]
  @test EC[i].mag.scale == true_msc[i]
end

@test EC[2].ot==true_ot
for i = 1:2
  @test ≈(EC[i].loc.lat, true_loc[1,i])
  @test ≈(EC[i].loc.lon, true_loc[2,i])
  @test ≈(EC[i].loc.dep, true_loc[3,i])
end

printstyled("    file read 2\n", color=:light_green)
H, R = read_qml(xml_evfile2)
H = H[1]
R = R[1]

# Check basic headers
@test H.typ == "earthquake"
@test H.id == "16461282"

# Check that the correct magnitude is retained
@test H.mag.val ≥ 9.0f0
@test H.mag.scale == "MW"

# Check H.loc
@test H.loc.lat ≈ 38.2963
@test H.loc.lon ≈ 142.498
@test H.loc.dep ≈ 19.7152
@test H.loc.nst == 2643
@test H.loc.gap ≈ 6.314
@test H.loc.dt ≈ 0.31
@test H.loc.se ≈ 2.1567
@test H.loc.dmin ≈ 0.917
@test H.loc.dmax ≈ 163.253
@test H.loc.src == "smi:ISC/origid=602227159,ISC"

# Check source params
@test R.id == "600002952"
@test R.m0 == 5.312e22
@test R.mt == [1.73e22, -2.81e21, -1.45e22, 2.12e22, 4.55e22, -6.57e21]
@test R.dm == [6.0e19, 5.0e19, 5.0e19, 6.8e20, 6.5e20, 4.0e19]
@test R.pax == [295.0 115.0 205.0; 55.0 35.0 0.0; 5.305e22 -5.319e22 1.4e20]
@test R.planes == [25.0 203.0; 80.0 10.0; 90.0 88.0]
@test R.st.dur == 70.0
@test R.misc["methodID"] == "smi:ISC/methodID=Best_double_couple"
@test R.misc["pax_desc"] == "azimuth, plunge, length"
@test split(R.src, ",")[2] == "GCMT"
@test R.misc["planes_desc"] == "strike, dip, rake"
@test R.misc["derivedOriginID"] == "600126955"

printstyled("    file write\n", color=:light_green)
xf = "test.xml"
if isfile(xf)
  safe_rm(xf)
end
write_qml(xf, H, R)
H1, R1 = read_qml(xf)
H1 = H1[1]
R1 = R1[1]
compare_SeisHdr(H1, H)
compare_SeisSrc(R1, R)

write_qml(xf, [H], [R])
write_qml(xf, H)
write_qml(xf, [H])
H1, R1 = read_qml(xf)
H1 = H1[4]
R1 = R1[2]
compare_SeisHdr(H1, H)
compare_SeisSrc(R1, R)

xstr = read(xml_evfile1)
io = open(xf, "w")
write(io, xstr)
close(io)
write_qml(xf, [H], [R])
H1, R1 = read_qml(xf)
H1 = H1[2:end]
R1 = R1[2:end]
H, R = read_qml(xml_evfile1)
for i in 1:length(R)
  compare_SeisHdr(H1[i], H[i])
  compare_SeisSrc(R1[i], R[i])
end

printstyled("      does trying to append a non-XML file error?\n", color=:light_green)
io = open(xf, "w")
write(io, rand(UInt8, 64))
close(io)
@test_throws ErrorException write_qml(xf, H, R)

# Clean up
safe_rm(xf)


text:
using quantumoptics
using multimode
using Optim

# Particles
const E0 = 0.2
const Nparticles = 1
const Nlevels = 12
const potential = multimode.BoxPotential
const particles = multimode.Bosons(Nparticles, Nlevels, E0, potential)

# Cavity mode
const index1 = 3
const cavitymodes1 = 12
const delta1 = -2.
const eta1 = 4.0
const kappa1 = 1.
const U0_1 = 0
const mode1 = multimode.CavityMode(index1, cavitymodes1, delta1, eta1, kappa1, U0_1)

# Multimode system
const scaling = 1./3
const system = multimode.MultimodeSystem(scaling, particles, [mode1])

# Bases
basis_particles = multimode.basis(particles)
basis_mode1 = multimode.basis(mode1)

# Initial state
const ψ₀ = basis_ket(basis_particles, 1) ⊗ basis_ket(basis_mode1, 1)
const ρ₀ = ψ₀ ⊗ dagger(ψ₀)

function fmin(x)
    N = length(x)//2
    y = x[1:N] + 1im*x[N+1:end]
    y /= norm(y,2)
    return multimode.multimode_steadystate.Hfunctional(system, Ket(basis_particles, y))
end
y0 = basis_ket(basis_particles, 1).data
N = length(y0)
x0 = [real(y0),imag(y0)]
x = Optim.optimize(fmin, x0).minimum
y = x[1:N] + 1im*x[N+1:end]
y /= norm(y,2)
x1_opt = Ket(basis_particles, y)
ρg1_opt = x1_opt ⊗ dagger(x1_opt)

function fmin2(x)
    N = length(x)//2
    y = x[1:N] + 1im*x[N+1:end]
    y /= norm(y,2)
    y = Ket(basis_particles, y)
    y_ = y - (dagger(x1_opt)*y)*x1_opt
    y_ /= sqrt(dagger(y_)*y_)
    return multimode.multimode_steadystate.Hfunctional(system, y_)
end
y0 = basis_ket(basis_particles, 1).data
N = length(y0)
x0 = [real(y0),imag(y0)]
x = Optim.optimize(fmin2, x0).minimum
y = x[1:N] + 1im*x[N+1:end]
y /= norm(y,2)
y = Ket(basis_particles, y)
y_ = y - (dagger(x1_opt)*y)*x1_opt
n = sqrt(dagger(y_)*y_)
y_ /= n
println("Norm: ", dagger(y_)*y_)
x2_opt = y_
ρg2_opt = x2_opt ⊗ dagger(x2_opt)

function fmin3(x)
    N = length(x)//2
    y = x[1:N] + 1im*x[N+1:end]
    y /= norm(y,2)
    y = Ket(basis_particles, y)
    y_ = y - (dagger(x1_opt)*y)*x1_opt - (dagger(x2_opt)*y)*x2_opt
    y_ /= sqrt(dagger(y_)*y_)
    return multimode.multimode_steadystate.Hfunctional(system, y_)
end
y0 = basis_ket(basis_particles, 1).data
N = length(y0)
x0 = [real(y0),imag(y0)]
x = Optim.optimize(fmin3, x0).minimum
y = x[1:N] + 1im*x[N+1:end]
y /= norm(y,2)
y = Ket(basis_particles, y)
y_ = y - (dagger(x1_opt)*y)*x1_opt - (dagger(x2_opt)*y)*x2_opt
n = sqrt(dagger(y_)*y_)
y_ /= n
println("Norm: ", dagger(y_)*y_)
x3_opt = y_
ρg3_opt = x3_opt ⊗ dagger(x3_opt)

T = [0.,100.0]

tout, ρt = multimode.timeevolution_master(system, T, ρ₀)

ρp = ptrace(ρt[end], [2])
ρf = ptrace(ρt[end], [1])
ρf /= trace(ρf)

U, S, V = svd(ρt[end].data)
println(S[1:5])

x1 = Ket(ρ₀.basis_l, U[:,1])
x2 = Ket(ρ₀.basis_l, U[:,2])
x3 = Ket(ρ₀.basis_l, U[:,3])
x4 = Ket(ρ₀.basis_l, U[:,4])
x5 = Ket(ρ₀.basis_l, U[:,5])

ρg1 = x1 ⊗ dagger(x1)
ρg2 = x2 ⊗ dagger(x2)
ρg3 = x3 ⊗ dagger(x3)
ρg4 = x4 ⊗ dagger(x4)
ρg5 = x5 ⊗ dagger(x5)

ρf1 = ptrace(ρg1, [1])
ρf2 = ptrace(ρg2, [1])
ρf3 = ptrace(ρg3, [1])
ρf4 = ptrace(ρg4, [1])
ρf5 = ptrace(ρg5, [1])

function f1(x)
    state = coherent_state(basis_mode1, complex(x[1], x[2]))
    return tracedistance(state ⊗ dagger(state), ρf1)
end
x = Optim.optimize(f1, Float64[0,0]).minimum
α1 = x[1] + 1im*x[2]

function f2(x)
    state = coherent_state(basis_mode1, complex(x[1], x[2]))
    return tracedistance(state ⊗ dagger(state), ρf2)
end
x = Optim.optimize(f2, Float64[0,0]).minimum
α2 = x[1] + 1im*x[2]

function f3(x)
    state = coherent_state(basis_mode1, complex(x[1], x[2]))
    return tracedistance(state ⊗ dagger(state), ρf3)
end
x = Optim.optimize(f3, Float64[0,0]).minimum
α3 = x[1] + 1im*x[2]

function f4(x)
    state = coherent_state(basis_mode1, complex(x[1], x[2]))
    return tracedistance(state ⊗ dagger(state), ρf4)
end
x = Optim.optimize(f4, Float64[0,0]).minimum
α4 = x[1] + 1im*x[2]

function f5(x)
    state = coherent_state(basis_mode1, complex(x[1], x[2]))
    return tracedistance(state ⊗ dagger(state), ρf5)
end
x = Optim.optimize(f5, Float64[0,0]).minimum
α5 = x[1] + 1im*x[2]

println("α₁ = ", α1)
println("α₂ = ", α2)
println("α₃ = ", α3)
println("α₄ = ", α4)
println("α₅ = ", α5)

Halpha1 = multimode.multimode_steadystate.Halpha(system, α1)
Halpha2 = multimode.multimode_steadystate.Halpha(system, α2)
Halpha3 = multimode.multimode_steadystate.Halpha(system, α3)

B_ = multimode.quantum.B(1, system)
D, V = eig(Hermitian(Halpha1.data))
x1_ = Ket(basis_particles, vec(V[:,1]))
x2_ = Ket(basis_particles, vec(V[:,2]))
x3_ = Ket(basis_particles, vec(V[:,3]))
x4_ = Ket(basis_particles, vec(V[:,4]))
x5_ = Ket(basis_particles, vec(V[:,5]))

println("α₁ = ", eta1*expect(B_, x1_)/(delta1+1im*kappa1))
println("α₂ = ", eta1*expect(B_, x2_)/(delta1+1im*kappa1))
println("α₃ = ", eta1*expect(B_, x3_)/(delta1+1im*kappa1))
println("α₄ = ", eta1*expect(B_, x4_)/(delta1+1im*kappa1))
println("α₅ = ", eta1*expect(B_, x5_)/(delta1+1im*kappa1))

println("α₁opt = ", eta1*expect(B_, x1_opt)/(delta1+1im*kappa1))
println("α₂opt = ", eta1*expect(B_, x2_opt)/(delta1+1im*kappa1))
println("α₃opt = ", eta1*expect(B_, x3_opt)/(delta1+1im*kappa1))

ρg1_ = x1_ ⊗ dagger(x1_)
ρg2_ = x2_ ⊗ dagger(x2_)
ρg3_ = x3_ ⊗ dagger(x3_)
ρg4_ = x4_ ⊗ dagger(x4_)
ρg5_ = x5_ ⊗ dagger(x5_)

x = linspace(-1,1,100)

n1 = multimode.particledensity(particles, x, ptrace(ρg1, [2]))
n2 = multimode.particledensity(particles, x, ptrace(ρg2, [2]))
n3 = multimode.particledensity(particles, x, ptrace(ρg3, [2]))
n4 = multimode.particledensity(particles, x, ptrace(ρg4, [2]))
n5 = multimode.particledensity(particles, x, ptrace(ρg5, [2]))

n1_ = multimode.particledensity(particles, x, ρg1_)
n2_ = multimode.particledensity(particles, x, ρg2_)
n3_ = multimode.particledensity(particles, x, ρg3_)
n4_ = multimode.particledensity(particles, x, ρg4_)
n5_ = multimode.particledensity(particles, x, ρg5_)

n1_opt = multimode.particledensity(particles, x, ρg1_opt)
n2_opt = multimode.particledensity(particles, x, ρg2_opt)
n3_opt = multimode.particledensity(particles, x, ρg3_opt)

using PyCall
@pyimport matplotlib.pyplot as plt
plt.figure()
plt.plot(x, n1, "k--")
plt.plot(x, n1_, "b")
plt.plot(x, n1_opt, "b--")
plt.plot(x, n2, "k--")
plt.plot(x, n2_opt, "g--")
plt.plot(x, n3_opt, "r--")
plt.plot(x, n2_, "g")
# plt.plot(x, n3, "k--")
plt.plot(x, n3_, "r")
# plt.plot(x, n4, "k--")
# plt.plot(x, n4_, "c")
# plt.plot(x, n5, "k--")
# plt.plot(x, n5_, "orange")
plt.show()
@assert false

tout, ρt = multimode.timeevolution_master(system, T, ρg1)
U, S, V = svd(ρt[end].data)

println(S[1:5])

ρg = S[1]*ρg1 + S[2]*ρg2

x = linspace(-1,1,100)
n1 = multimode.particledensity(particles, x, ptrace(ρg1, [2]))
n1 = multimode.particledensity(particles, x, ptrace(ρg1, [2]))
n2 = multimode.particledensity(particles, x, ptrace(ρg2, [2]))
n3 = multimode.particledensity(particles, x, ptrace(ρg3, [2]))
n4 = multimode.particledensity(particles, x, ptrace(ρg4, [2]))
n5 = multimode.particledensity(particles, x, ptrace(ρg5, [2]))

using PyCall
@pyimport matplotlib.pyplot as plt

q = qfunc(ptrace(ρg1, [1]), [-4:0.1:4], [-4:0.1:4])
plt.imshow(q)

plt.figure()
q = qfunc(ptrace(ρg2, [1]), [-4:0.1:4], [-4:0.1:4])
plt.imshow(q)

plt.figure()
q = qfunc(ptrace(ρg3, [1]), [-4:0.1:4], [-4:0.1:4])
plt.imshow(q)

plt.figure()
q = qfunc(ptrace(ρg4, [1]), [-4:0.1:4], [-4:0.1:4])
plt.imshow(q)

plt.figure()
q = qfunc(ptrace(ρg5, [1]), [-4:0.1:4], [-4:0.1:4])
plt.imshow(q)

# plt.figure()
# q = qfunc(ρf-rho_f_try, [-4:0.1:4], [-4:0.1:4])
# plt.imshow(q)
# plt.colorbar()

plt.figure()
plt.plot(x, n1)
plt.plot(x, n2)
plt.plot(x, n3)
plt.plot(x, n4)
plt.plot(x, n5)
plt.show()

# Hp = multimode.Hamiltonian(particle)
# B = multimode.B(1, system)

# ξ = 7.

# Hplus = Hp + ξ*B
# Hminus = Hp - ξ*B

# Dplus, Vplus = eig(Hplus.data)
# Dminus, Vminus = eig(Hminus.data)

# vplus1 = Ket(particle.basis, Vplus[:,1])
# vminus1 = Ket(particle.basis, Vminus[:,1])
# vplus2 = Ket(particle.basis, Vplus[:,2])
# vminus2 = Ket(particle.basis, Vminus[:,2])
# vplus3 = Ket(particle.basis, Vplus[:,3])
# vminus3 = Ket(particle.basis, Vminus[:,3])
# vplus4 = Ket(particle.basis, Vplus[:,4])
# vminus4 = Ket(particle.basis, Vminus[:,4])

# #vplus = sqrt(0.8)*vplus1 + sqrt(0.15)*vplus2 + sqrt(0.05)*vplus4
# #vminus = sqrt(0.8)*vminus1 + sqrt(0.15)*vminus2 + sqrt(0.05)*vminus4
# #vplus = sqrt(0.98)*vplus1 + sqrt(0.02)*vplus2
# #vminus = sqrt(0.98)*vminus1 + sqrt(0.02)*vminus2
# vplus = vplus1
# vminus = vminus1

# n_xplus = particledensity(particle, [-1:0.01:1], vplus)
# n_xminus = particledensity(particle, [-1:0.01:1], vminus)

# exp_B = dagger(vplus)*B*vplus
# η = sqrt((delta1^2 + kappa1^2)/(delta1*exp_B)*ξ)
# α = η*exp_B/(delta1 + 1im*kappa1)

# println("<B> = ", exp_B)
# println("η = ", η)
# println("α = ", α)
# println("n = ", abs2(α))

# # for i=4:5
# #     vplus = Ket(particle.basis, Vplus[:,i])
# #     vminus = Ket(particle.basis, Vminus[:,i])
# #     n_xplus = particledensity(particle, [-1:0.01:1], vplus)
# #     n_xminus = particledensity(particle, [-1:0.01:1], vminus)
# #     plt.plot((n_xplus + n_xminus)/2)
# # end

# # plt.show()
# # @assert false
# plt.plot(n_xplus)
# plt.plot(n_xminus)
# plt.figure()
# plt.plot((n_xplus + n_xminus)/2)
# #plt.plot(n_xplus)

# mode1.eta = η
# ρplus₀ = tensor(vplus, dagger(vplus))
# ρminus₀ = tensor(vminus, dagger(vminus))
# Ψfplus₀ = coherent_state(mode1.basis, α)
# Ψfminus₀ = coherent_state(mode1.basis, -α)
# #Ψf₀ = coherent_state(mode1.basis, 0)
# #println(expect(destroy(mode1.basis), Ψf₀))
# ρfplus₀ = tensor(Ψfplus₀, dagger(Ψfplus₀))
# ρfminus₀ = tensor(Ψfminus₀, dagger(Ψfminus₀))
# ρ₀ = 0.5*tensor(ρplus₀, ρfplus₀) + 0.5*tensor(ρminus₀, ρfminus₀)

# #println(expect(destroy(mode1.basis), ρf₀))
# #@assert false
# #ρ₀ = tensor(basis_ket(system.basis, 1), basis_bra(system.basis, 1))
# T, ρ_t = multimode.evolve_master(system, [0.:1:5.], ρ₀)
# exp_n = expect(number(mode1.basis), ptrace(ρ_t[end], [1]))
# println("<a^t a> = ", exp_n)
# n2 = particledensity(particle, [-1:0.01:1], ptrace(ρ_t[2], [2]))
# n3 = particledensity(particle, [-1:0.01:1], ptrace(ρ_t[3], [2]))
# n4 = particledensity(particle, [-1:0.01:1], ptrace(ρ_t[4], [2]))
# n = particledensity(particle, [-1:0.01:1], ptrace(ρ_t[end], [2]))
# q = qfunc(ptrace(ρ_t[end], [1]), [-3:0.1:3], [-3:0.1:3])

# using PyCall
# @pyimport matplotlib.pyplot as plt
# plt.plot(n2)
# plt.plot(n3)
# plt.plot(n4)
# plt.plot(n)
# plt.figure()
# plt.imshow(q)
# plt.show()
# #println([calc_exp_B(ρ) for ρ in ρ_t])
# #println(real(diag(ptrace(ρ_t[1], [2]).data)))
# #println(real(diag(ptrace(ρ_t[end], [2]).data)))

# #println(norm((Hplus*v1-D[1]*v1).data))


text:
using Yao.YaoBlocks.Optimise: simplify
using ProgressMeter
export simplify
export neighboring_pairs, heisenberg1D, train!, tfim

neighboring_pairs(n) = [i=>mod1(i+1, n) for i in 1:n]

YaoExtensions.variational_circuit(n, nlayers) = variational_circuit(n, nlayers, neighboring_pairs(n))

tfim(n) = sum(put(n, i=>Z) * put(n, mod1(i+1, n)=>Z) for i in 1:n) + sum(put(n, i=>X) for i in 1:n)

_heisenberg_S(n, i) = sum(put(n, i=>σ) * put(n, i+1=>σ) for σ in (X, Y, Z))
heisenberg1D(n) = simplify(sum(_heisenberg_S(n, i) for i in 1:n-1))

using Flux.Optimise

function train!(opt, epochs, ham, circuit; verbose=true)
    history = Float64[]
    push!(history, 1000)

    n = nqubits(ham)
    early_stop = 10
    @showprogress for k in 1:epochs
        # the expectation is calculated on a complex matrix
        # thus we just use the real part here
        E = expect(ham, zero_state(n)=>circuit) |> real

        if verbose
            @info "step=$k, E/n=$(E/4n)"
        end

        # early stop
        if length(history) > 1 && abs(history[end] - E/4n) < 1e-6
            @info "Early stopping: $early_stop"
            early_stop -= 1
        end
        if early_stop < 1
            @info "Early stopping!"
            return history
        end
        push!(history, E/4n)
        _, grad = expect'(ham, zero_state(n)=>circuit)
        ps = parameters(circuit)
        Optimise.update!(opt, ps, grad)
        popdispatch!(circuit, ps)

    end
    return history
end


text:
function Base.write(io::IO, records::Vector{Record}) #Note: we assume the indexes have been bumpped and the open ends are correct.
    for record in records
        Base.write(io, record, '\n')
    end
    return position(io)
end

function Base.write(io::IO, record::Record)
    # delim = '\t'
    delim = ' '
    strings = convert(Vector{String}, record)

    Base.join(io, strings, delim)

    return position(io)
end


text:
lattice = load_lattice(joinpath(@__DIR__, "GaAs20.json"))
atoms   = load_atoms(  joinpath(@__DIR__, "GaAs20.json"))
atoms   = attach_pseudos(atoms, Ga="hgh/pbe/ga-q3.hgh", As="hgh/pbe/as-q5.hgh")

cases = standard_cases()
cases["adaptive_03"] = (damping=DFTK.AdaptiveDamping(;α_trial_min=0.3, α_min=0.05), tol=1e-12)
cases["adaptive_04"] = (damping=DFTK.AdaptiveDamping(;α_trial_min=0.4, α_min=0.05), tol=1e-12)
cases["adaptive_05"] = (damping=DFTK.AdaptiveDamping(;α_trial_min=0.5, α_min=0.05), tol=1e-12)

function load_GaAs_20_potential(basis)
    guessfile = joinpath(@__DIR__, "GaAs20_initV.jld2")
    if !isfile(guessfile)
        ρ0 = guess_density(basis)
        _, ham = DFTK.energy_hamiltonian(basis, nothing, nothing; ρ=ρ0)
        V = DFTK.total_local_potential(ham)
        V .+= 3.0rand(Float64, size(V)...)

        jldopen(guessfile, "w") do file
            file["V"] = V
        end
    end

    jldopen(guessfile, "r") do file
        file["V"]
    end
end

# Load guess from file
config[@__DIR__] = (
    directory        = @__DIR__,
    lattice          = lattice,
    atoms            = atoms,
    magnetic_moments = [],
    #
    kwargs_model     = (smearing=Smearing.None(), temperature=0.0),
    Ecut             = 20,
    kgrid            = [4, 4, 1],
    supersampling    = 1.5,
    potential        = load_GaAs_20_potential,
    kwargs_scf       = (; ),
    mixing           = SimpleMixing(),
    cases            = cases,
)

max_stars_repo_name: mfherbst/supporting-adaptive-damping

max_issues_repo_name: mfherbst/supporting-adaptive-damping

max_forks_repo_name: mfherbst/supporting-adaptive-damping


text:
# Low-level wrappers for Portaudio calls

# General type aliases
const PaTime = Cdouble
const PaError = Cint
const PaSampleFormat = Culong
const PaDeviceIndex = Cint
const PaHostApiIndex = Cint
const PaHostApiTypeId = Cint
# PaStream is always used as an opaque type, so we're always dealing
# with the pointer
const PaStream = Ptr{Void}
const PaStreamCallback = Void
const PaStreamFlags = Culong

const paNoFlag = PaStreamFlags(0x00)

const PA_NO_ERROR = 0
const PA_INPUT_OVERFLOWED = -10000 + 19
const PA_OUTPUT_UNDERFLOWED = -10000 + 20

# sample format types
const paFloat32 = PaSampleFormat(0x01)
const paInt32   = PaSampleFormat(0x02)
const paInt24   = PaSampleFormat(0x04)
const paInt16   = PaSampleFormat(0x08)
const paInt8    = PaSampleFormat(0x10)
const paUInt8   = PaSampleFormat(0x20)
const paNonInterleaved = PaSampleFormat(0x80000000)

const type_to_fmt = Dict{Type, PaSampleFormat}(
    Float32 => 1,
    Int32   => 2,
    # Int24   => 4,
    Int16   => 8,
    Int8    => 16,
    UInt8   => 3
)

const PaStreamCallbackResult = Cint
# Callback return values
const paContinue = PaStreamCallbackResult(0)
const paComplete = PaStreamCallbackResult(1)
const paAbort = PaStreamCallbackResult(2)

function Pa_Initialize()
    err = ccall((:Pa_Initialize, libportaudio), PaError, ())
    handle_status(err)
end

function Pa_Terminate()
    err = ccall((:Pa_Terminate, libportaudio), PaError, ())
    handle_status(err)
end

Pa_GetVersion() = ccall((:Pa_GetVersion, libportaudio), Cint, ())

function Pa_GetVersionText()
    versionPtr = ccall((:Pa_GetVersionText, libportaudio), Ptr{Cchar}, ())
    unsafe_string(versionPtr)
end

# Host API Functions

# A Host API is the top-level of the PortAudio hierarchy. Each host API has a
# unique type ID that tells you which native backend it is (JACK, ALSA, ASIO,
# etc.). On a given system you can identify each backend by its index, which
# will range between 0 and Pa_GetHostApiCount() - 1. You can enumerate through
# all the host APIs on the system by iterating through those values.

# PaHostApiTypeId values
const pa_host_api_names = Dict{PaHostApiTypeId, String}(
    0 => "In Development", # use while developing support for a new host API
    1 => "Direct Sound",
    2 => "MME",
    3 => "ASIO",
    4 => "Sound Manager",
    5 => "Core Audio",
    7 => "OSS",
    8 => "ALSA",
    9 => "AL",
    10 => "BeOS",
    11 => "WDMKS",
    12 => "Jack",
    13 => "WASAPI",
    14 => "AudioScience HPI"
)

type PaHostApiInfo
    struct_version::Cint
    api_type::PaHostApiTypeId
    name::Ptr{Cchar}
    deviceCount::Cint
    defaultInputDevice::PaDeviceIndex
    defaultOutputDevice::PaDeviceIndex
end

Pa_GetHostApiInfo(i) = unsafe_load(ccall((:Pa_GetHostApiInfo, libportaudio),
                                   Ptr{PaHostApiInfo}, (PaHostApiIndex,), i))

# Device Functions

type PaDeviceInfo
    struct_version::Cint
    name::Ptr{Cchar}
    host_api::PaHostApiIndex
    max_input_channels::Cint
    max_output_channels::Cint
    default_low_input_latency::PaTime
    default_low_output_latency::PaTime
    default_high_input_latency::PaTime
    default_high_output_latency::PaTime
    default_sample_rate::Cdouble
end

Pa_GetDeviceCount() = ccall((:Pa_GetDeviceCount, libportaudio), PaDeviceIndex, ())

Pa_GetDeviceInfo(i) = unsafe_load(ccall((:Pa_GetDeviceInfo, libportaudio),
                                 Ptr{PaDeviceInfo}, (PaDeviceIndex,), i))

Pa_GetDefaultInputDevice() = ccall((:Pa_GetDefaultInputDevice, libportaudio),
                                   PaDeviceIndex, ())

Pa_GetDefaultOutputDevice() = ccall((:Pa_GetDefaultOutputDevice, libportaudio),
                                    PaDeviceIndex, ())

# Stream Functions

type Pa_StreamParameters
    device::PaDeviceIndex
    channelCount::Cint
    sampleFormat::PaSampleFormat
    suggestedLatency::PaTime
    hostAPISpecificStreamInfo::Ptr{Void}
end

type PaStreamInfo
    structVersion::Cint
    inputLatency::PaTime
    outputLatency::PaTime
    sampleRate::Cdouble
end

# function Pa_OpenDefaultStream(inChannels, outChannels,
#                               sampleFormat::PaSampleFormat,
#                               sampleRate, framesPerBuffer)
#     streamPtr = Ref{PaStream}(0)
#     err = ccall((:Pa_OpenDefaultStream, libportaudio),
#                 PaError, (Ref{PaStream}, Cint, Cint,
#                           PaSampleFormat, Cdouble, Culong,
#                           Ref{Void}, Ref{Void}),
#                 streamPtr, inChannels, outChannels, sampleFormat, sampleRate,
#                 framesPerBuffer, C_NULL, C_NULL)
#     handle_status(err)
#
#     streamPtr[]
# end
#
function Pa_OpenStream(inParams, outParams,
                       sampleRate, framesPerBuffer,
                       flags::PaStreamFlags,
                       callback, userdata)
    streamPtr = Ref{PaStream}(0)
    err = ccall((:Pa_OpenStream, libportaudio), PaError,
                (Ref{PaStream},
                Ptr{Pa_StreamParameters},
                Ptr{Pa_StreamParameters},
                Cdouble, Culong, PaStreamFlags,
                Ptr{Void}, Ptr{Void}),
                streamPtr,
                inParams, outParams,
                sampleRate, framesPerBuffer, flags,
                callback, userdata)
    handle_status(err)
    streamPtr[]
end

function Pa_StartStream(stream::PaStream)
    err = ccall((:Pa_StartStream, libportaudio), PaError,
                (PaStream,), stream)
    handle_status(err)
end

function Pa_StopStream(stream::PaStream)
    err = ccall((:Pa_StopStream, libportaudio), PaError,
                (PaStream,), stream)
    handle_status(err)
end

function Pa_CloseStream(stream::PaStream)
    err = ccall((:Pa_CloseStream, libportaudio), PaError,
                (PaStream,), stream)
    handle_status(err)
end

function Pa_GetStreamReadAvailable(stream::PaStream)
    avail = ccall((:Pa_GetStreamReadAvailable, libportaudio), Clong,
                (PaStream,), stream)
    avail >= 0 || handle_status(avail)
    avail
end

function Pa_GetStreamWriteAvailable(stream::PaStream)
    avail = ccall((:Pa_GetStreamWriteAvailable, libportaudio), Clong,
                (PaStream,), stream)
    avail >= 0 || handle_status(avail)
    avail
end

function Pa_ReadStream(stream::PaStream, buf::Array, frames::Integer=length(buf),
                       show_warnings::Bool=true)
    frames <= length(buf) || error("Need a buffer at least $frames long")
    err = ccall((:Pa_ReadStream, libportaudio), PaError,
                (PaStream, Ptr{Void}, Culong),
                stream, buf, frames)
    handle_status(err, show_warnings)
    buf
end

function Pa_WriteStream(stream::PaStream, buf::Array, frames::Integer=length(buf),
                        show_warnings::Bool=true)
    frames <= length(buf) || error("Need a buffer at least $frames long")
    err = ccall((:Pa_WriteStream, libportaudio), PaError,
                (PaStream, Ptr{Void}, Culong),
                stream, buf, frames)
    handle_status(err, show_warnings)
    nothing
end

# function Pa_GetStreamInfo(stream::PaStream)
#     infoptr = ccall((:Pa_GetStreamInfo, libportaudio), Ptr{PaStreamInfo},
#             (PaStream, ), stream)
#
#     unsafe_load(infoptr)
# end
#
# General utility function to handle the status from the Pa_* functions
function handle_status(err::PaError, show_warnings::Bool=true)
    if err == PA_OUTPUT_UNDERFLOWED || err == PA_INPUT_OVERFLOWED
        if show_warnings
            msg = ccall((:Pa_GetErrorText, libportaudio),
                        Ptr{Cchar}, (PaError,), err)
            warn("libportaudio: " * unsafe_string(msg))
        end
    elseif err != PA_NO_ERROR
        msg = ccall((:Pa_GetErrorText, libportaudio),
                    Ptr{Cchar}, (PaError,), err)
        error("libportaudio: " * unsafe_string(msg))
    end
end