text:
#!/usr/bin/env julia
# File              : src/VSFlow.jl
# Author            : Denis Dumoulin <denis.dumoulin@uclouvain.be>
# Date              : 30.08.2021
# Last Modified Date: 28.09.2021
module VSFlow

using LinearAlgebra
using Printf
using Logging
using Plots; gr()

include("VortexElement.jl")
include("Airfoil.jl")
include("History.jl")

export uniform, heavepitch, circularmotion
export naca00, gaw1, circle, ellipse
export Profile, profilerun
export History, getimpulse, getboundpotential, plotaero, plotcps, plotϕs

"""
PROFILE object
"""
mutable struct Profile
	profileshape::Function
	N::Int64
	panels::Array{LinearPanel, 1}
	V::Float64
	vortex_points::Array{VortexPoint, 1}
	constant_panel::ConstantPanel
	old_α::Float64
	ω::Float64
	γs::Vector{Float64}
	γg::Float64
	oldγg::Float64
	average_length::Float64
	dt::Float64
	T::Float64
	δ::Float64
	ϵ::Float64
	timelapse::Float64
	is_init::Bool
	θg::Float64
	θ0::Float64
	β::Float64
	θ1::Float64
	θ2::Float64
	old_panel::ConstantPanel
	pivot_location::Array{Float64, 1}
	Ut::Float64
	Vt::Float64
	force_control::Array{Float64, 1}
	trailing::Array{Float64, 1}
	fname::String
	Ainv::Array{Float64, 2}

	#Lumping
	η::Float64
	τ::Int64
	Tmin::Int64
	v_corr::Vector{Float64}
	sheet_size::Int64
	last_vp_ind::Int64
    islumped::Bool

	profID::String
    history::History
    isgaussian::Bool
end

"""
    Profile(; id, profileshape::Function, x0, ẋ0, N, dt, T, δ = 1e-2, ϵ = 1e-2, (η, Tmin, Smin) = zeros(3))

Constructs a `Profile` object.
If `eta`, `Tmin`, `Smin` are all `0`, then no lumping operation occurs.

# Keyword Arguments
 - `id`: `String` used to identify the profile.
 - `profileshape`: `function` determining the profile shape.
 - `N`: number of panels on the surface of the profile.
 - `x0`: initial position `[X0, Y0, α0]` of the quarter chord of the profile.
 - `ẋ0`: initial velocity `[Ẋ0, Ẏ0, α̇0]` of the quarter chord of the profile.
 - `dt`: timestep of the simulation.
 - `T`: horizon time.
 - `isgaussian`: `` determining if the vortex kernel is Gaussian or Low-order Algebraic.
 - `δ = 1e-2`: kernel cut-off width.
 - `ϵ = 1e-2`: percentage of the average panel length to chop the trailing edge.
 - `η = 0`: maximum error due to lumping.
 - `Tmin = 0`: minimum time between two successive active vortices.
 - `Smin = 0`: minimum vortex sheet length in the wake.

"""
function Profile(; id, profileshape::Function, x0, ẋ0, N, dt, T,
                    isgaussian = false,
                    δ = 1e-2,
                    ϵ = 1e-2,
                    lumpargs = zeros(3))
	iseven(N) || error("Need an even number of panels!")
    islumped = (lumpargs != zeros(3))
    η, Tmin, sheet_size = lumpargs

	x0[1] -= .25
	panels = genprofile(profileshape, N, id, x0[1:end-1])
	γs = zeros(N + 1)
	average_length = sum(2 .* map(p->p.b, panels))/N

    vortex_points = []
	fname = id*"_np"*string(N)*"_dt"*string(dt)*"_T"*string(T)*"_dv"*string(δ)*"_eps"*string(ϵ)
	islumped && (fname = fname*"_lump_errMax"*string(η)*"_Tmin"*string(Tmin)*"_ss"*string(sheet_size))
	fname = replace(fname, "."=>"")

	pivot_location = [x0[1] + .25, x0[2]]
	V = sum(map(pa->getbodyvolume(pa, pivot_location...), panels))
	for p in panels
		setangle!(p, x0[end], pivot_location)
	end
    U0, V0, α̇0 = ẋ0

	θ0 = panels[1].θ - panels[end].θ + sign(panels[end].θ)*pi
	θ0 = abs(atan(sin(θ0), cos(θ0)))
	θ1 = .5θ0
	θ2 = .5θ0
	θg = panels[1].θ - θ1

	X1 = .5(panels[1].X1+panels[end].X2)
	Y1 = .5(panels[1].Y1+panels[end].Y2)
	trailing = [X1, Y1]

	#Constant panel positioning
	b = dt - ϵ*average_length
	X2 = X1 + 2b*cos(θg)
	Y2 = Y1 - 2b*sin(θg)
	constant_panel = BuildPanel(ConstantPanel, X1, Y1, X2, Y2, id)
	cropedge!(constant_panel, ϵ*average_length, false)
	cropedge!(panels[1], ϵ*average_length, false)
	cropedge!(panels[end], ϵ*average_length, true)
    h = History(N, dt, T)
	return Profile(profileshape, N, panels, V, vortex_points, constant_panel, x0[end],
				   -α̇0, γs, 0, 0, average_length, dt, T, δ, ϵ, 0,
				   true, θg, θ0, 0, θ1, θ2, constant_panel, pivot_location, U0, V0,
				   zeros(3), trailing, fname, zeros(N+2, N+2),
				   η, 0, Tmin, zeros(2), sheet_size, 1, islumped, id, h, isgaussian)
end

function copier!(p::Profile, p2::Profile)
	p.profileshape = p2.profileshape
	p.N = p2.N
	p.panels = p2.panels
	p.V = p2.V
	p.vortex_points = p2.vortex_points
	p.constant_panel = p2.constant_panel
	p.old_α = p2.old_α
	p.ω = p2.ω
	p.γs = p2.γs
	p.γg = p2.γg
	p.oldγg = p2.oldγg
	p.average_length = p2.average_length
	p.dt = p2.dt
	p.T = p2.T
	p.δ = p2.δ
	p.ϵ = p2.ϵ
	p.timelapse = p2.timelapse
	p.is_init = p2.is_init
	p.θg = p2.θg
	p.θ0 = p2.θ0
    p.β = p2.β
	p.θ1 = p2.θ1
	p.θ2 = p2.θ2
	p.old_panel = p2.old_panel
	p.pivot_location = p2.pivot_location
	p.Ut = p2.Ut
	p.Vt = p2.Vt
	p.force_control = p2.force_control
	p.trailing = p2.trailing
	p.fname = p2.fname
	p.Ainv = p2.Ainv
	p.η = p2.η
	p.τ = p2.τ
	p.Tmin = p2.Tmin
	p.v_corr = p2.v_corr
	p.sheet_size = p2.sheet_size
	p.last_vp_ind = p2.last_vp_ind
    p.islumped = p2.islumped
	p.profID = p2.profID
	p.history = p2.history
	p.isgaussian = p2.isgaussian
	return nothing
end

"""
    genprofile(profile, N, position, args...)

Generate a profile with a wedged trailing edge.
"""
function genprofile(profileshape::Function, N::Int64, profID, position::Array)
	θ = range(0, pi, length=N÷2 + 1)
	x = @. .5cos(θ) + .5
    ye, yi = profileshape(x)

	L = []
	for i in 1:length(x)-1
		@inbounds	push!(L, BuildPanel(LinearPanel, x[i] + position[1],
										  ye[i] + position[2], x[i+1] + position[1], ye[i+1] + position[2], profID))
	end
	for i in length(x)-1:-1:1
		@inbounds push!(L, BuildPanel(LinearPanel, x[i+1] + position[1],
										 yi[i+1] + position[2], x[i] + position[1], yi[i] + position[2], profID))
	end
	return L
end

"""
    getprofilelengths(p::Profile)

Return an array containing lengths of all panels.
"""
function getprofilelengths(p::Profile)
	λ = 2map(pr->pr.b, p.panels)
	pushfirst!(λ, 0)
	return cumsum(λ)
end

"""
    inbodyframe(p::Profile, X, Y)

Return the vector `[X, Y]` in a frame attached to the pivot point of the body.
"""
function inbodyframe(p::Profile, X, Y)
	X0, Y0 = p.pivot_location
	x, y = rotate(X-X0, Y-Y0, p.old_α)
	return [x, y]
end

"""
    getparamsinrefframe(panel, X, Y)

Return the velocity parameters from `panel` at point `X`, `Y` in the inertial frame.
"""
function getparamsinrefframe(panel::LinearPanel, X, Y)
	u, v = getparams(panel, X, Y)
	uv1 = rotate(u[1], v[1], -(π+panel.θ))
	uv2 = rotate(u[2], v[2], -(π+panel.θ))
	U = [uv1[1], uv2[1]]
	V = [uv1[2], uv2[2]]
	return [U, V]
end

"""
    getparamsinrefframe(panel, X, Y)

Return the velocity parameters from `panel` at point `X`, `Y` in the inertial frame.
"""
function getparamsinrefframe(panel::ConstantPanel, X, Y)
	u, v = getparams(panel, X, Y)
	return rotate(u, v, -(π+panel.θ))
end

"""
    getboundpanelsinducedvel(profile, X, Y)

Return the velocity at point `X`, `Y` due to the body panels.
"""
function getboundpanelsinducedvel(p::Profile, X, Y)
	X0, Y0 = p.pivot_location
	U = V = 0
	for (i, panel) in enumerate(p.panels)
		u, v = getparamsinrefframe(panel, X, Y)
		@inbounds U += u[1]*p.γs[i] + u[2]*p.γs[i+1]
		@inbounds V += v[1]*p.γs[i] + v[2]*p.γs[i+1]
	end
	int = sum(map((pa)->getmotionparams(pa, X, Y, p.ω, p.Ut, p.Vt, X0, Y0), p.panels))

	return [U, V] + int
end

"""
    motionupdate!(p::Profile, X, Ẋ)

Update the position, AoA and velocities of the body.
"""
function motionupdate!(p::Profile, X, Ẋ)
	X0, Y0 = p.pivot_location
	Xt, Yt, θ = X
	α = -θ
	p.Ut, p.Vt, p.ω = Ẋ
	for panel in p.panels
		translate!(panel, Xt-X0, Yt-Y0)
		setangle!(panel, α-p.old_α, [Xt, Yt])
	end
	p.old_α = α
	p.pivot_location = [Xt, Yt]
	X1 = .5(p.panels[1].X1 + p.panels[end].X2)
	Y1 = .5(p.panels[1].Y1 + p.panels[end].Y2)
	p.trailing = [X1, Y1]
	return nothing
end

"""
    globalvelocity(p::Profile, XY, incloud=false)

Return the velocity at point `XY`.

# Keyword Arguments
 - `incloud = false`: flag indicating whether `XY` is part of the cloud of point vortices.
"""
function globalvelocity(p::Profile, XY, incloud=false)
	UV = getboundpanelsinducedvel(p, XY[1], XY[2])
	UV += p.γg*getparamsinrefframe(p.constant_panel, XY[1], XY[2])
	!incloud && !isempty(p.vortex_points) &&
        (UV += sum(map(vp->getinducedvelocity(vp, XY[1], XY[2], p.δ, p.isgaussian), p.vortex_points)))
	return UV
end

"""
    setforcecontrol!(p::Profile, fx, fy, mz)

Update forces acting on the body pivot point.
"""
function setforcecontrol!(p::Profile, fx, fy, mz)
	p.force_control = [fx, fy, -mz]
	return nothing
end

"""
    gettrailingedgevel!(p::Profile)

Return the fluid velocity at the trailing edge expressed in the inertial frame.
"""
function gettrailingedgevel!(p::Profile)
	v1 = p.γs[1]
	v2 = -p.γs[end]
    p.β = atan(tan(.5p.θ0)*(v2-v1)/(v2+v1))
    p.θ1 = .5p.θ0 + p.β
    p.θ2 = .5p.θ0 - p.β
	p.θg = p.panels[1].θ - p.θ1
    u3 = .5*(v1*cos(p.θ1) + v2*cos(p.θ2))

	return [-u3*cos(p.θg), u3*sin(p.θg)]
end

"""
    placeconstantpanel!(p::Profile, X2, Y2)

Place the uniform panel at location `X2`, `Y2`.
"""
function placeconstantpanel!(p::Profile, X2, Y2)
	X1 = p.trailing[1] + p.ϵ*p.average_length*cos(p.θg)
	Y1 = p.trailing[2] - p.ϵ*p.average_length*sin(p.θg)
	p.constant_panel = BuildPanel(ConstantPanel, X1, Y1, X2, Y2, p.profID)
	return nothing
end

"""
    step!(p::Profile)

Update the simulation in time.

#Keyword Arguments
 - `is4thorder = true`: indicates if a RK4 or a ForwardEuler scheme is used.
 - `need_reset = true`: indicates whether the uniform panel needs to be turned
    into a point vortex.
"""
function step!(p::Profile, is4thorder=true, need_reset=true)
	X0, Y0 = p.pivot_location

	X1 = p.trailing[1] + p.ϵ*p.average_length*cos(p.θg)
	Y1 = p.trailing[2] - p.ϵ*p.average_length*sin(p.θg)

	XY0 = map(vp->[vp.X, vp.Y], p.vortex_points)
	XA0 = [X0, Y0, -p.old_α]
	ẊA0 = [p.Ut, p.Vt, p.ω]
	XP0 = [X1, Y1]

	order = 1 + is4thorder*3
	β = [0, .5, .5, 1]
	ζ = is4thorder*[.5, 1., 1., .5]/3. + !is4thorder*[1., 0, 0, 0]
	dXY = XY0
	dXA = XA0
	dXP = XP0
	dẊA = ẊA0
	XY = XY0
	XA = XA0
	XP = XP0
	ẊA = ẊA0
	XYL = XY0
	XAL = XA0
	XPL = XP0
	ẊAL = ẊA0
	for k in 1:order
		if k > 1
			@inbounds XYL = XY0 + β[k]*dXY
			@inbounds XAL = XA0 + β[k]*dXA
			@inbounds XPL = XP0 + β[k]*dXP
			@inbounds ẊAL = ẊA0 + β[k]*dẊA
			for (i, vp) in enumerate(p.vortex_points)
				@inbounds vp.X = XYL[i][1]
				@inbounds vp.Y = XYL[i][2]
			end
			@inbounds motionupdate!(p, XAL, ẊAL)
			placeconstantpanel!(p, XPL...)
			setγs!(p)
		end
		dXY = p.dt*(getcloudvelocities(p.vortex_points, p.δ, p.isgaussian) .+
                        map(x->globalvelocity(p, x, true), XYL))

		!isempty(p.vortex_points) && (dXY[p.last_vp_ind] += p.vortex_points[p.last_vp_ind].dv*p.dt)
		dXP = p.dt*gettrailingedgevel!(p)
		dẊA = p.dt*p.force_control
		dXA = p.dt*ẊAL
		@inbounds XY += dXY*ζ[k]
		@inbounds XP += dXP*ζ[k]
		@inbounds ẊA += dẊA*ζ[k]
		@inbounds XA += dXA*ζ[k]
	end

	for (i, vp) in enumerate(p.vortex_points)
		@inbounds vp.X = XY[i][1]
		@inbounds vp.Y = XY[i][2]
	end
	motionupdate!(p, XA, ẊA)
	placeconstantpanel!(p, XP...)
	setγs!(p)
	gettrailingedgevel!(p)

	need_reset && resetconstantpanel!(p)

	p.timelapse += p.dt
	p.τ += 1
	return nothing
end

"""
    function resetconstantpanel!(p)

Transform the constant panel into a point vortex.
Then recompute the vortex strength distribution on the body.
"""
function resetconstantpanel!(p)
	X = .5(p.constant_panel.X1+p.constant_panel.X2)
	Y = .5(p.constant_panel.Y1+p.constant_panel.Y2)
	Γ = 2p.γg*p.constant_panel.b

	if Γ != 0
		new_vp = VortexPoint(X, Y, Γ)
		push!(p.vortex_points, new_vp)
	end

	X1 = p.trailing[1] + p.ϵ*p.average_length*cos(p.θg)
	Y1 = p.trailing[2] - p.ϵ*p.average_length*sin(p.θg)
    p.old_panel = deepcopy(p.constant_panel)
	placeconstantpanel!(p, X1, Y1)

    p.oldγg = p.γg
	p.γg = 0
	setγs!(p)
	return nothing
end

"""
    lumpvortices!(p::Profile)

Lump the last shed point vortex into the active vortex if lumping criteria are met.
Used instead of `step!` if lumping is desired.
"""
function lumpvortices!(p::Profile)
	if p.last_vp_ind > length(p.vortex_points) - p.sheet_size + 1
		step!(p)
		p.τ = 0
		return 0
	end

	t_index = p.last_vp_ind
	s_index = p.last_vp_ind + 1
	if sign(p.vortex_points[t_index].Γ) != sign(p.vortex_points[s_index].Γ)
		p.τ = 0
		p.vortex_points[t_index].dv = zeros(2)
		p.last_vp_ind += 1
		step!(p)
		return 0
	end

	γ̇ = p.vortex_points[s_index].Γ/p.dt
	s_vp = p.vortex_points[s_index]
	t_vp = p.vortex_points[t_index]
	s_imp = getvorteximpulse(p, s_vp)
	t_imp = getvorteximpulse(p, t_vp)

	grads = zeros(p.N+1, 2)
	for (i, panel) in enumerate(p.panels)
		@inbounds grads[i, :] = getvelocitygradient(t_vp, panel.Xc, panel.Yc, π+panel.θ, p.δ)
	end
	∇γ = (p.Ainv*grads)'

	κ = map(pa->[pa.Y1, -pa.X1], p.panels)
	push!(κ, [p.panels[end].Y2, -p.panels[end].X2])
	κ = hcat(κ...)
	A = mapslices(fook->[mapslices(foog->trapz(fook.*foog, getprofilelengths(p)), ∇γ, dims=2)], κ, dims=2)

	∇p = [0 1; -1 0] + hcat(A...)'
	p.v_corr = (∇p \ (s_imp - t_imp))γ̇/t_vp.Γ
	η = getforceerror!(p, s_index, t_index)
	return η
end

"""
    getforceerror!(p::Profile, s_index, t_index)

Return the error in impulse due to the lumping of vortex with index `s_index`.
"""
function getforceerror!(p::Profile, s_index, t_index)
	no_transfer_p = deepcopy(p)
	no_transfer_p.τ = 0
	no_transfer_p.vortex_points[t_index].dv = zeros(2)
	no_transfer_p.last_vp_ind += 1
	step!(no_transfer_p)
	no_transfer_imp = getimpulse(no_transfer_p)[1:2]

	transfer!(p, s_index, t_index)
	step!(p)
	imp = getimpulse(p)[1:2]

	η = norm(imp - no_transfer_imp)/p.dt
	if η > p.η && p.τ >= p.Tmin
		@info "No Lumping allowed"
		copier!(p, no_transfer_p)
		η = 0.
	end
	no_transfer_p = nothing
	return η
end

"""
    getvorteximpulse(p, vp)

Return the impulse linked with vortex `vp` and its image on the body.
"""
function getvorteximpulse(p::Profile, vp)
	b = zeros(p.N + 1)
	for (m, panel) in enumerate(p.panels)
		U, V = getinducedvelocity(vp, panel.Xc, panel.Yc, p.δ, p.isgaussian)
		_, V = rotate(U, V, π+panel.θ)
		@inbounds b[m] = -V
	end
	b[end] = -vp.Γ
	γs = p.Ainv * b

	X0, Y0 = p.hisory.X[1], p.hisory.X[2]
	Xp, Yp = p.pivot_location
	panels_impulse = sum(map((pa, γL, γR)->getimpulse(pa, γL, γR, p.ω, p.Ut, p.Vt, X0, Y0, Xp, Yp, false)[1:2], p.panels, γs[1:end-1], γs[2:end]))
	return [vp.Y-Y0, X0-vp.X] + panels_impulse/vp.Γ
end

"""
    transfer!(p, s_index, t_index)

Lump vortex of index `s_index` into active vortex of index `t_index`.
"""
function transfer!(p::Profile, s_index, t_index)
	p.vortex_points[t_index].Γ += p.vortex_points[s_index].Γ
	p.vortex_points[t_index].dv = p.v_corr
	deleteat!(p.vortex_points, s_index)
	return nothing
end

"""
    setγs!(p::Profile)

Compute the vortex strength distribution around the profile.
"""
function setγs!(p::Profile)
	X0, Y0 = p.pivot_location
	A = zeros(p.N + 2, p.N + 2)
	b = zeros(p.N + 2)

	m::UInt64 = 1
	for panel in p.panels

		V = - .5rotate(p.Ut - p.ω*(panel.Yc-Y0), p.Vt + p.ω*(panel.Xc-X0), π+panel.θ)[2]
		for vp in p.vortex_points
			Up, Vp = getinducedvelocity(vp, panel.Xc, panel.Yc, p.δ, p.isgaussian)
			Up, Vp = rotate(Up, Vp, π+panel.θ)
			V += Vp
		end
		Ub, Vb = sum(map((pa)->getmotionparams(pa, panel.Xc, panel.Yc, p.ω, p.Ut, p.Vt, X0, Y0), p.panels))

		@inbounds b[m] -= V + rotate(Ub, Vb, π+panel.θ)[2]

		n::UInt64 = 1
		for other in p.panels
			v1, v2 = getlinearspeedcoeffs(panel, other, true)
			@inbounds A[m, n] += v1
			@inbounds A[m, n+1] += v2
			n += 1
		end
		v = getconstantspeedcoeffs(panel, p.constant_panel, true)
		@inbounds A[m, n+1] = v

		#Kelvin
		@inbounds A[p.N+1, m] += panel.b
		@inbounds A[p.N+1, m+1] += panel.b
		m += 1
	end
	@inbounds A[p.N+1, m+1] += 2p.constant_panel.b

	int = -sum(map(pa->getmotionkelvin(pa, p.ω, p.Ut, p.Vt, X0, Y0), p.panels))
	@inbounds b[p.N+1] = -sum(map(vp->vp.Γ, p.vortex_points)) + int

	#Unsteady Kutta
	@inbounds A[p.N+2, 1] = cos(p.θ1)
	@inbounds A[p.N+2, p.N+1] = cos(p.θ2)
	@inbounds A[p.N+2, p.N+2] = -1

	if p.constant_panel.b == 0
		@inbounds p.Ainv = inv(A[1:p.N+1, 1:p.N+1])
		@inbounds p.γs = p.Ainv * b[1:p.N+1]
	else
		p.Ainv = inv(A)
		p.γs = p.Ainv * b
	end

	if length(p.γs) == p.N + 2
		p.γg = p.γs[end]
		pop!(p.γs)
	end

	p.is_init = false
	return nothing
end

"""
    getimpulse(p::Profile)

Return the linear/angular impulse exerted by the fluid on the body.
"""
function getimpulse(p::Profile)
	X0, Y0 = p.history.X[1, 1], p.history.X[1, 2]
	Xp, Yp = p.pivot_location
	panels_impulse = sum(map((pa, γL, γR)->getimpulse(pa, γL, γR, p.ω, p.Ut, p.Vt, X0, Y0, Xp, Yp, true), p.panels, p.γs[1:end-1], p.γs[2:end]))

	if !isempty(p.vortex_points)
		points_impulse = sum(map(vp->vp.Γ*[vp.Y-Y0, -vp.X+X0, -.5norm2(vp.Y - Y0, vp.X - X0)], p.vortex_points))
	else
		points_impulse = zeros(3)
	end
	return panels_impulse + points_impulse
end

"""
    setcoeffimpulse!(p::Profile, idx, bdorder)

Set aerodynamic coefficients computed with impulse conservation in the flow.
"""
function setcoeffimpulse!(p::Profile, idx, bdorder)
	X0, Y0 = p.history.X[1, 1], p.history.X[1, 2]
	xyz_imp = getimpulse(p::Profile)
    oldimpulse = getlastvalues(p.history, :P, idx-1, bdorder)

	cd, cl, cm = -2backwarddifference(xyz_imp, oldimpulse, p.dt)
	cm -= dot(p.pivot_location - [X0, Y0], [cl, -cd])

    p.history.P[idx, :] = xyz_imp
    p.history.aci[idx, :] = [cd, cl, cm]
end

"""
    setcoeffnoca!(p::Profile, idx, bdorder)

Set aerodynamic coefficients computed with a control volume approach.
"""
function setcoeffnoca!(p::Profile, idx, bdorder)
	X0, Y0 = p.history.X[1, 1], p.history.X[1, 2]
    Xp, Yp = p.pivot_location

	panels_impulse = sum(map((pa, γL, γR)->getimpulse(pa, γL, γR, p.ω, p.Ut, p.Vt, Xp, Yp, Xp, Yp, true), p.panels, p.γs[1:end-1], p.γs[2:end]))
    oldnocaimpulse = getlastvalues(p.history, :NP, idx-1, bdorder)

	u = 2p.old_panel.b/p.dt
	energyflux = sum(map((pa, γL, γR)->getenergyflux(pa, γL, γR, p.ω, 0, 0, Xp, Yp, Xp, Yp), p.panels, p.γs[1:end-1], p.γs[2:end]))
	vorticityflux = p.oldγg*u*([p.trailing[2]-Yp, Xp - p.trailing[1], -.5norm2.((p.trailing - p.pivot_location)...)])
	Δm = sum(map(pa->getangularimpulsecorrection(pa, p.ω, 0, 0, p.Ut, p.Vt, Xp, Yp, Xp, Yp), p.panels))

	f = energyflux - backwarddifference(panels_impulse, oldnocaimpulse, p.dt) - vorticityflux
	cd, cl, cm = 2f

    p.history.NP[idx, :] = panels_impulse
    p.history.acn[idx, :] = [cd, cl, cm + Δm]
end

"""
    getboundpotential(p::Profile)

Integrates the velocity on the profile to obtain the velocity potential `ϕ`.
"""
function getboundpotential(p::Profile)
	X0, Y0 = p.pivot_location
	ϕs = zeros(p.N+1)
	tmp = 0
	for i in 1:p.N
		@inbounds tmp += getpotential(p.panels[i], p.γs[i], p.γs[i+1], p.ω, p.Ut, p.Vt, X0, Y0)
		@inbounds ϕs[i+1] = tmp
	end
	return ϕs
end

"""
    setcoeffpotential!(p::Profile, idx, bdorder)

Set aerodynamic coefficients using pressure integration.
"""
function setcoeffpotential!(p::Profile, idx, bdorder)
	ϕs = getboundpotential(p)
    oldϕs = getlastvalues(p.history, :ϕs, idx-1, bdorder)

	X0, Y0 = p.pivot_location
	Us = map(pa->p.Ut - p.ω*(pa.Y1 - Y0), p.panels)
	Vs = map(pa->p.Vt + p.ω*(pa.X1 - X0), p.panels)
	push!(Us, p.Ut - p.ω*(p.panels[end].Y2 - Y0))
	push!(Vs, p.Vt + p.ω*(p.panels[end].X2 - X0))

	cps = 1. .+ norm2.(Us, Vs) .- p.γs.^2 .- 2backwarddifference(ϕs, oldϕs, p.dt)
	cd, cl, cm = sum(map((pa, p1, p2)->getpressureforce(pa, p1, p2, p.pivot_location...), p.panels, cps[1:end-1], cps[2:end]))

    p.history.ϕs[idx, :] = ϕs
    p.history.cps[idx, :] = cps
    p.history.acp[idx, :] = [cd, cl, cm]
end

"""
    profilerun(p::Profile, accfunc; animate=false, is4thorder=true, bdorder=4)

Simulate the testcase.

# Arguments
 - `p`: `Profile` simulation to run.
 - `accfunc`: `Function` giving the force applied on the pivot point of the body.

# Keyword Arguments
 - `animate = false`: `Bool` indicating whether an animation has to be generated.
 - `is4thorder = true`: `Bool` indicating if the order of the Runge-Kutta integrator is 4.
                        If `false`, then a Forward Euler method is used.
 - `bdorder = 4`: Order of the backward difference schemes for time derivatives.
"""
function profilerun(p::Profile, accfunc; animate=false, is4thorder=true, bdorder=4)
	run(`echo "==================================================================="`)
	run(pipeline(`git log`, `head -6`))
	run(`echo "==================================================================="`)
	println(" Params:")
	println("---------")
	println("Panel Number     :"*string(p.N))
	println("Timestep         :"*string(p.dt))
	println("Horizon Time     :"*string(p.T))
	println("Blob δ           :"*string(p.δ))
	println("-----------------------------------------------------------")

	setγs!(p)

	count = 0
	if animate
		anim = Animation()
	end
    while(count <= 1 + Int(p.T/p.dt))
		if count%10 == 0
			@info @sprintf "%.3f/%.3f" p.timelapse p.T
		end

		η = 0.
		setforcecontrol!(p, accfunc(p.timelapse)[3]...)
		if p.islumped
			η, t, bytes, _, _ = @timed lumpvortices!(p)
		else
			_, t, bytes, _, _ = @timed step!(p, is4thorder)
		end
		count += 1

        ##Update history
        θ = p.β/p.θ0
		Γb = sum(map(vp->vp.Γ, p.vortex_points))
		setcoeffpotential!(p, count, bdorder)
        setcoeffimpulse!(p, count, bdorder)
        setcoeffnoca!(p, count, bdorder)
        updatehistory!(p.history, count, p.timelapse,
                       [p.pivot_location[1], p.pivot_location[2], p.old_α],
                       [p.Ut, p.Vt, p.ω],  Γb, θ)

		##Animation
		skip=1
		if animate && count%skip == 0
			x = map(vp->vp.X, p.vortex_points)
			y = map(vp->vp.Y, p.vortex_points)
			Γ = map(vp->vp.Γ, p.vortex_points)
            traj = plot(p.history.X[1:count, 1], p.history.X[1:count, 2],
                        line=(1, :lightgray), aspect_ratio=1, xlim=(-.5-p.T, 1.2),
                        ylim=(-3, 3), framestyle=:none, grid=false, ticks=false,
                        colorbar=false, dpi=400)
			pp = scatter!(x, y, zcolor=Γ, markercolor=:coolwarm, markersize=3,
                          markerstrokewidth=0, clims=(-.2,.2))
			xp = map(panel->panel.X1, p.panels)
			yp = map(panel->panel.Y1, p.panels)
			push!(xp, p.panels[end].X2)
			push!(yp, p.panels[end].Y2)
			plot!(pp, xp, yp, line=(2, :black), fill=(0, :black), legend=false)

			frame(anim)
		end
	end

	if animate
		gif(anim, p.fname*".gif", fps=20)
        println("Animation saved!")
	end

	println("===================================================================")
	println("DONE")
	println("===================================================================")
end

end


text:
import ClimaCore:
    Fields,
    Domains,
    Topologies,
    Meshes,
    DataLayouts,
    Operators,
    Geometry,
    Spaces

using OrdinaryDiffEq: ODEProblem, solve, SSPRK33

import Logging
import TerminalLoggers
Logging.global_logger(TerminalLoggers.TerminalLogger())

const FT = Float64

n = 32
z₀ = FT(0)
z₁ = FT(10)
t₀ = FT(0)
t₁ = FT(10)
μ = FT(-1 / 2)
ν = FT(5)
𝓌 = FT(1)
δ = FT(1)

domain = Domains.IntervalDomain(
    Geometry.ZPoint{FT}(z₀),
    Geometry.ZPoint{FT}(z₁),
    boundary_names = (:bottom, :top),
)
mesh = Meshes.IntervalMesh(domain, nelems = n)

cs = Spaces.CenterFiniteDifferenceSpace(mesh)
fs = Spaces.FaceFiniteDifferenceSpace(cs)
zc = Fields.coordinate_field(cs)
zp = (z₀ + z₁ / n / 2):(z₁ / n):(z₁ - z₁ / n / 2)

function gaussian(z, t; μ = -1 // 2, ν = 1, 𝓌 = 1, δ = 1)
    return exp(-(z - μ - 𝓌 * t)^2 / (4 * ν * (t + δ))) / sqrt(1 + t / δ)
end
function ∇gaussian(z, t; μ = -1 // 2, ν = 1, 𝓌 = 1, δ = 1)
    return -2 * (z - μ - 𝓌 * t) / (4 * ν * (δ + t)) *
           exp(-(z - μ - 𝓌 * t)^2 / (4 * ν * (δ + t))) / sqrt(1 + t / δ)
end

T = gaussian.(zc.z, -0; μ = μ, δ = δ, ν = ν, 𝓌 = 𝓌)
V = Geometry.WVector.(ones(FT, fs))

# Solve Adv-Diff Equation: ∂_t T = α ∇²T
z₀ = FT(0)
z₁ = FT(10)

function ∑tendencies!(dT, T, z, t)

    ic2f = Operators.InterpolateC2F()
    bc_vb = Operators.SetValue(FT(gaussian(z₀, t; ν = ν, δ = δ, 𝓌 = 𝓌, μ = μ)))
    bc_vt = Operators.SetValue(FT(gaussian(z₁, t; ν = ν, δ = δ, 𝓌 = 𝓌, μ = μ)))
    bc_gb = Operators.SetGradient(
        Geometry.WVector(FT(∇gaussian(z₀, t; ν = ν, δ = δ, 𝓌 = 𝓌, μ = μ))),
    )
    bc_gt = Operators.SetGradient(
        Geometry.WVector(FT(∇gaussian(z₁, t; ν = ν, δ = δ, 𝓌 = 𝓌, μ = μ))),
    )

    #   Upwind Biased Product
    #   UB = Operators.UpwindBiasedProductC2F(
    #       bottom = Operators.Extrapolate(),
    #       top = bc_vt,
    #   )
    #   ∂ = Operators.GradientF2C()
    #   return @. dT = -∂(UB(V, ic2f(T)))

    A = Operators.AdvectionC2C(bottom = bc_vb, top = Operators.Extrapolate())

    gradc2f = Operators.GradientC2F(bottom = bc_vb, top = bc_gt)
    divf2c = Operators.DivergenceF2C()

    return @. dT = divf2c(ν * gradc2f(T)) - A(V, T)
end

@show ∑tendencies!(similar(T), T, nothing, 0.0);

# Solve the ODE operator
Δt = 0.0001

prob = ODEProblem(∑tendencies!, T, (t₀, t₁))
sol = solve(
    prob,
    SSPRK33(),
    dt = Δt,
    saveat = 10000 * Δt,
    progress = true,
    progress_message = (dt, u, p, t) -> t,
);

ENV["GKSwstype"] = "nul"
using ClimaCorePlots, Plots
Plots.GRBackend()

dirname = "advect_diffusion"
path = joinpath(@__DIR__, "output", dirname)
mkpath(path)

anim = Plots.@animate for (nt, u) in enumerate(sol.u)
    Plots.plot(
        u,
        xlim = (0, 1),
        ylim = (-1, 10),
        title = "$(nt-1) s",
        lc = :black,
        lw = 2,
        ls = :dash,
        label = "Approx Sol.",
        legend = :outerright,
        m = :o,
        xlabel = "T(z)",
        ylabel = "z",
    )
    Plots.plot!(
        gaussian.(zp, nt - 1; μ = μ, δ = δ, ν = ν, 𝓌 = 𝓌),
        zp,
        xlim = (0, 1),
        ylim = (-1, 10),
        title = "$(nt) s",
        lc = :red,
        lw = 2,
        label = "Analytical Sol.",
        m = :x,
    )
end
Plots.mp4(anim, joinpath(path, "advect_diffusion.mp4"), fps = 10)
Plots.png(
    Plots.plot(sol.u[end], xlim = (0, 1)),
    joinpath(path, "advect_diffusion_end.png"),
)

function linkfig(figpath, alt = "")
    # buildkite-agent upload figpath
    # link figure in logs if we are running on CI
    if get(ENV, "BUILDKITE", "") == "true"
        artifact_url = "artifact://$figpath"
        print("\033]1338;url='$(artifact_url)';alt='$(alt)'\a\n")
    end
end

linkfig(
    relpath(
        joinpath(path, "advect_diffusion_end.png"),
        joinpath(@__DIR__, "../.."),
    ),
    "Advection-Diffusion End Simulation",
)

ext: jl


text:
# This file includes all functions that serve to create data structures that
# record partition information and relevant network parameters

# NOTE: this file is updated to incorporate the case when soem lines are inactivated while parsed by PowerModels.jl
#       however, we are still assuming that all buses are active (which might cause errors if not true)

function find_neighbor_buses(data::Dict{String, Any}, N_g::Set{Int64})::Set{Int64}
    neighbors = Set{Int64}()
    for line in keys(data["branch"])
        if data["branch"][line]["br_status"] != PM.pm_component_status_inactive["branch"] &&
                (data["branch"][line]["f_bus"] in N_g || data["branch"][line]["t_bus"] in N_g)
            push!(neighbors, data["branch"][line]["f_bus"], data["branch"][line]["t_bus"])
        end
    end
    return setdiff!(neighbors, N_g)
end

function generate_subnet_data(data::Dict{String, Any}, N_g::Vector{Int64})
    return generate_subnet_data(data, Set(N_g))
end

function generate_subnet_data(data::Dict{String, Any}, N_g::Set{Int64})
    sub_data = deepcopy(data)
    nbus = find_neighbor_buses(data, N_g)
    N = Set(parse.(Int, collect(keys(data["bus"]))))
    # deactivate unrelated buses
    for i in setdiff!(N, N_g, nbus)
        sub_data["bus"]["$(i)"]["bus_type"] = PM.pm_component_status_inactive["bus"]
    end
    # deactivate lines between neighbor buses
    for line in keys(data["branch"])
        if sub_data["branch"][line]["f_bus"] in nbus && data["branch"][line]["t_bus"] in nbus
            sub_data["branch"][line]["br_status"] = PM.pm_component_status_inactive["branch"]
        end
    end
    # deactivate unrelated generators (for those at neighbor nodes)
    for (_, gen) in sub_data["gen"]
        if !(gen["gen_bus"] in N_g)
            gen["gen_status"] = PM.pm_component_status_inactive["gen"]
        end
    end
    sub_data["cut_bus"] = nbus
    return sub_data
end

function ref_add_cut_bus!(ref::Dict{Symbol, Any}, data::Dict{String, Any})
    for (nw, nw_ref) in ref[:it][:pm][:nw]
        nw_ref[:cut_bus] = Dict(i => data["bus"]["$(i)"] for i in data["cut_bus"])
    end
end

function ref_add_cut_branch!(ref::Dict{Symbol, Any}, data::Dict{String, Any})
    for (nw, nw_ref) in ref[:it][:pm][:nw]
        # set up cut branches and arcs
        nw_ref[:cut_branch] = Dict(parse(Int, x.first) => x.second for x in data["branch"] if
            x.second["br_status"] != PM.pm_component_status_inactive["branch"] &&
            (x.second["f_bus"] in keys(nw_ref[:bus])) + (x.second["t_bus"] in keys(nw_ref[:bus])) +
            (x.second["f_bus"] in keys(nw_ref[:cut_bus])) + (x.second["t_bus"] in keys(nw_ref[:cut_bus])) == 3)
        nw_ref[:cut_arcs_from] = [(i,branch["f_bus"],branch["t_bus"]) for (i,branch) in nw_ref[:cut_branch]]
        nw_ref[:cut_arcs_to]   = [(i,branch["t_bus"],branch["f_bus"]) for (i,branch) in nw_ref[:cut_branch]]
        nw_ref[:cut_arcs] = [nw_ref[:cut_arcs_from]; nw_ref[:cut_arcs_to]]
    end
end

#=
function ref_add_cut_bus_arcs_refs!(ref::Dict{Symbol, Any}, data::Dict{String, Any})
    for (nw, nw_ref) in ref[:it][:pm][:nw]
        cut_bus_arcs = Dict((i, Tuple{Int,Int,Int}[]) for (i,bus) in merge(nw_ref[:bus], nw_ref[:cut_bus]))
        for (l,i,j) in nw_ref[:cut_arcs]
            push!(cut_bus_arcs[i], (l,i,j))
        end
        nw_ref[:cut_bus_arcs] = cut_bus_arcs
    end
end
=#

function ref_add_global_bus!(ref::Dict{Symbol, Any}, data::Dict{String, Any})
    function calc_all_buspair_parameters(buses, branches, conductor_ids, ismulticondcutor)
        bus_lookup = Dict(bus["index"] => bus for (i,bus) in buses)
        branch_lookup = Dict(branch["index"] => branch for (i,branch) in branches)
        buspair_indexes = Set((branch["f_bus"], branch["t_bus"]) for (i,branch) in branch_lookup)
        bp_branch = Dict((bp, typemax(Int64)) for bp in buspair_indexes)

        if ismulticondcutor
            bp_angmin = Dict((bp, [-Inf for c in conductor_ids]) for bp in buspair_indexes)
            bp_angmax = Dict((bp, [ Inf for c in conductor_ids]) for bp in buspair_indexes)
        else
            @assert(length(conductor_ids) == 1)
            bp_angmin = Dict((bp, -Inf) for bp in buspair_indexes)
            bp_angmax = Dict((bp,  Inf) for bp in buspair_indexes)
        end

        for (l,branch) in branch_lookup
            i = branch["f_bus"]
            j = branch["t_bus"]

            if ismulticondcutor
                for c in conductor_ids
                    bp_angmin[(i,j)][c] = max(bp_angmin[(i,j)][c], branch["angmin"][c])
                    bp_angmax[(i,j)][c] = min(bp_angmax[(i,j)][c], branch["angmax"][c])
                end
            else
                bp_angmin[(i,j)] = max(bp_angmin[(i,j)], branch["angmin"])
                bp_angmax[(i,j)] = min(bp_angmax[(i,j)], branch["angmax"])
            end

            bp_branch[(i,j)] = min(bp_branch[(i,j)], l)
        end

        buspairs = Dict((i,j) => Dict(
            "branch"=>bp_branch[(i,j)],
            "angmin"=>bp_angmin[(i,j)],
            "angmax"=>bp_angmax[(i,j)],
            "tap"=>branch_lookup[bp_branch[(i,j)]]["tap"],
            "vm_fr_min"=>bus_lookup[i]["vmin"],
            "vm_fr_max"=>bus_lookup[i]["vmax"],
            "vm_to_min"=>bus_lookup[j]["vmin"],
            "vm_to_max"=>bus_lookup[j]["vmax"]
            ) for (i,j) in buspair_indexes
        )

        # add optional parameters
        for bp in buspair_indexes
            branch = branch_lookup[bp_branch[bp]]
            if haskey(branch, "rate_a")
                buspairs[bp]["rate_a"] = branch["rate_a"]
            end
            if haskey(branch, "c_rating_a")
                buspairs[bp]["c_rating_a"] = branch["c_rating_a"]
            end
        end

        return buspairs
    end
    for (nw, nw_ref) in ref[:it][:pm][:nw]
        nw_ref[:all_bus] = Dict(j["bus_i"] => j for (_, j) in data["bus"])
        nw_ref[:all_branch] = Dict{Int, Any}(parse(Int, k) => v for (k,v) in data["branch"] if v["br_status"] != PM.pm_component_status_inactive["branch"])
        nw_ref[:all_buspairs] = calc_all_buspair_parameters(nw_ref[:all_bus], nw_ref[:all_branch], nw_ref[:conductor_ids], haskey(nw_ref, :conductors))
    end
end


text:
"""
		getrange(n)

		returns block of indices for a particular thred
"""
function getrange(n)
    tid = Threads.threadid()
    nt = Threads.nthreads()
    d , r = divrem(n, nt)
    from = (tid - 1) * d + min(r, tid - 1) + 1
    to = from + d - 1 + (tid ≤ r ? 1 : 0)
    from:to
end

subsetof(bags) = bags[getrange(length(bags))]


text:
function add_velocity_bound!(model::AbstractGameModel, game_con::GameConstraintValues, v_max::AbstractVector, v_min::AbstractVector)
	n = model.n
	p = model.p
	@assert length(v_max) == length(v_min) == p
	for i = 1:p
		if (v_max[i] != Inf) || (v_min[i] != -Inf)
			add_velocity_bound!(model, game_con, i, v_max[i], v_min[i])
		end
	end
	return nothing
end

function add_velocity_bound!(model::AbstractGameModel, game_con::GameConstraintValues, i::Int, v_max::Real, v_min::Real)
	@assert (v_max != Inf) || (v_min != -Inf)
	n = model.n
	p = model.p
	# Compute the bound
	x_max = Inf*ones(n)
	x_min = -Inf*ones(n)
	vel_index_i = velocity_index(model, i)
	x_max[vel_index_i] = v_max
	x_min[vel_index_i] = v_min
	# Add the state bound to all players
	for j = 1:p
		add_state_bound!(game_con, j, x_max, x_min)
	end
	return nothing
end

function velocity_index(model::UnicycleGame, i::Int)
	@assert i ∈ (1:model.p)
	return model.pz[i][4]
end

function velocity_index(model::BicycleGame, i::Int)
	@assert i ∈ (1:model.p)
	return model.pz[i][3]
end

function velocity_index(model::DoubleIntegratorGame, i::Int)
	error("Velocity Index is not implemented for DoubleIntegratorGame.")
	return nothing
end


text:
# Generative model using the change of variables formula
# Author: Philipp Witte, pwitte3@gatech.edu
# Date: January 2020

using LinearAlgebra, InvertibleNetworks, JOLI, PyPlot, Flux, Random, Test, JLD, Statistics
import Flux.Optimise.update!
using Distributions: Uniform

# Random seed
Random.seed!(66)

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

# Load original data X (size of n1 x n2 x nc x ntrain)
X_orig = load("../../data/seismic_samples_32_by_32_num_10k.jld")["X"]
n1, n2, nc, nsamples = size(X_orig)
AN = ActNorm(nsamples)
X_orig = AN.forward(X_orig) # zero mean and unit std

# Split in training - testing
ntrain = Int(nsamples*.9)
ntest = nsamples - ntrain

# Dimensions after wavelet squeeze to increase no. of channels
nx = Int(n1)
ny = Int(n2)
n_in = Int(nc)

# Apply wavelet squeeze (change dimensions to -> n1/2 x n2/2 x nc*4 x ntrain)
X_train = zeros(Float32, nx, ny, n_in, ntrain)
for j=1:ntrain
    X_train[:, :, :, j:j] = X_orig[:, :, :, j:j]
end

X_test = zeros(Float32, nx, ny, n_in, ntest)
for j=1:ntest
    X_test[:, :, :, j:j] = X_orig[:, :, :, ntrain+j:ntrain+j]
end

# Create network
n_hidden = 32
batchsize = 4
L = 1
K = 8
CH = NetworkMultiScaleConditionalHINT(n_in, n_hidden, L, K; k1=3, k2=3, p1=1, p2=1, s1=1, s2=1, split_scales=false)
Params = get_params(CH)

# Data modeling function
function model_data(X, A)
    nx, ny, nc, nb = size(X)
    Y = reshape(A*reshape(X, :, nb), nx, ny, nc, nb)
    return Y
end

# Forward operator, precompute phase in ambient dimension
function phasecode(n)
	F = joDFT(n; DDT=Float32,RDT=ComplexF64)
	phase=F*(adjoint(F)*exp.(1im*2*pi*convert(Array{Float32},rand(dist,n))))
	phase = phase ./abs.(phase)
	sgn = sign.(convert(Array{Float32},randn(n)))
	# Return operator
	return M = joDiag(sgn) * adjoint(F) * joDiag(phase)*F
end

# Generate observed data
dist = Uniform(-1, 1)
input_dim = (n1, n2)
subsamp = 2
M = phasecode(prod(input_dim))
R = joRestriction(prod(input_dim),randperm(prod(input_dim))[1:Int(round(prod(input_dim)/subsamp))]; DDT=Float32, RDT=Float32);
A_flat = R*M;
A = A_flat'*A_flat

Y_train = model_data(X_train, A)
Y_test = model_data(X_test, A)

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

# Loss
function loss(CH, X, Y)
    Zx, Zy, logdet = CH.forward(X, Y)
    f = -log_likelihood(tensor_cat(Zx, Zy)) - logdet
    ΔZ = -∇log_likelihood(tensor_cat(Zx, Zy))
    ΔZx, ΔZy = tensor_split(ΔZ)
    ΔX, ΔY = CH.backward(ΔZx, ΔZy, Zx, Zy)[1:2]
    return f, ΔX, ΔY
end

# Training
maxiter = 1000
opt = Flux.ADAM(1f-3)
lr_step = 100
lr_decay_fn = Flux.ExpDecay(1f-3, .9, lr_step, 0.)
fval = zeros(Float32, maxiter)

for j=1:maxiter

    # Evaluate objective and gradients
    idx = randperm(ntrain)[1:batchsize]
    X = X_train[:, :, :, idx]
    Y = Y_train[:, :, :, idx]

    fval[j] = loss(CH, X, Y)[1]
    mod(j, 10) == 0 && (print("Iteration: ", j, "; f = ", fval[j], "\n"))

    # Update params
    for p in Params
        update!(opt, p.data, p.grad)
        update!(lr_decay_fn, p.data, p.grad)
    end
    clear_grad!(CH)
end

####################################################################################################
# Plotting

# Testing
test_size = 100
idx = randperm(ntest)[1:test_size]  # draw random samples from testing data
X = X_test[:, :, :, idx]
Y = Y_test[:, :, :, idx]
Zx_, Zy_ = CH.forward(X, Y)[1:2]

Zx = randn(Float32, size(Zx_))
Zy = randn(Float32, size(Zy_))
X_, Y_ = CH.inverse(Zx, Zy)

# Now select single fixed sample from all Ys
idx = 1
X_fixed = X[:, :, :, idx:idx]
Y_fixed = Y[:, :, :, idx:idx]
Zy_fixed = CH.forward_Y(Y_fixed)

# Draw new Zx, while keeping Zy fixed
CH.forward_Y(X) # set X dimensions in forward pass (this needs to be fixed)
X_post = CH.inverse(Zx, Zy_fixed.*ones(Float32, size(Zx_)))[1]

# Plot one sample from X and Y and their latent versions
figure(figsize=[8,8])
ax1 = subplot(2,2,1); imshow(X[:, :, 1, 1], cmap="gray", aspect="auto"); title(L"Model space: $x \sim \hat{p}_x$")
ax2 = subplot(2,2,2); imshow(Y[:, :, 1, 1], cmap="gray", aspect="auto"); title(L"Noisy data $y=x+n$ ")
ax3 = subplot(2,2,3); imshow(X_[:, :, 1, 1], cmap="gray", aspect="auto"); title(L"Model space: $x = f(zx|zy)^{-1}$")
ax4 = subplot(2,2,4); imshow(Y_[:, :, 1, 1], cmap="gray", aspect="auto"); title(L"Data space: $y = f(zx|zy)^{-1}$")

# Plot various samples from X and Y
figure(figsize=[16,8])
i = randperm(test_size)[1:4]
ax1 = subplot(2,4,1); imshow(X_[:, :, 1, i[1]], cmap="gray", aspect="auto"); title(L"Model space: $x = f(zx|zy)^{-1}$")
ax2 = subplot(2,4,2); imshow(X_[:, :, 1, i[2]], cmap="gray", aspect="auto"); title(L"Model space: $x = f(zx|zy)^{-1}$")
ax3 = subplot(2,4,3); imshow(X_[:, :, 1, i[3]], cmap="gray", aspect="auto"); title(L"Model space: $x = f(zx|zy)^{-1}$")
ax4 = subplot(2,4,4); imshow(X_[:, :, 1, i[4]], cmap="gray", aspect="auto"); title(L"Model space: $x = f(zx|zy)^{-1}$")
ax5 = subplot(2,4,5); imshow(X[:, :, 1, i[1]], cmap="gray", aspect="auto"); title(L"Model space: $x \sim \hat{p}_x$")
ax6 = subplot(2,4,6); imshow(X[:, :, 1, i[2]], cmap="gray", aspect="auto"); title(L"Model space: $x \sim \hat{p}_x$")
ax7 = subplot(2,4,7); imshow(X[:, :, 1, i[3]], cmap="gray", aspect="auto"); title(L"Model space: $x \sim \hat{p}_x$")
ax8 = subplot(2,4,8); imshow(X[:, :, 1, i[4]], cmap="gray", aspect="auto"); title(L"Model space: $x \sim \hat{p}_x$")

# Plot posterior samples, mean and standard deviation
figure(figsize=[16,8])
X_post_mean = mean(X_post; dims=4)
X_post_std = std(X_post; dims=4)
ax1 = subplot(2,4,1); imshow(X_fixed[:, :, 1, 1], cmap="gray", aspect="auto"); title("True x")
ax2 = subplot(2,4,2); imshow(X_post[:, :, 1, 1], cmap="gray", aspect="auto"); title(L"Post. sample: $x = f(zx|zy_{fix})^{-1}$")
ax3 = subplot(2,4,3); imshow(X_post[:, :, 1, 2], cmap="gray", aspect="auto"); title(L"Post. sample: $x = f(zx|zy_{fix})^{-1}$")
ax4 = subplot(2,4,4); imshow(X_post_mean[:, :, 1, 1], cmap="gray", aspect="auto"); title("Posterior mean")
ax5 = subplot(2,4,5); imshow(Y_fixed[:, :, 1, 1], cmap="gray", aspect="auto");  title(L"Noisy data $y_i=x_i+n$ ")
ax6 = subplot(2,4,6); imshow(X_post[:, :, 1, 4], cmap="gray", aspect="auto"); title(L"Post. sample: $x = f(zx|zy_{fix})^{-1}$")
ax7 = subplot(2,4,7); imshow(X_post[:, :, 1, 5], cmap="gray", aspect="auto"); title(L"Post. sample: $x = f(zx|zy_{fix})^{-1}$")
ax8 = subplot(2,4,8); imshow(X_post_std[:, :, 1,1], cmap="binary", aspect="auto", vmin=0, vmax=0.9*maximum(X_post_std));
colorbar(); title("Posterior std");


text:
"""
    uncond_QFI_transverse(t, N, κ, ω)

Analytical unconditional QFI for ``N`` atoms interacting with independent transverse
noise with coupling ``\\κ``, in a magnetic field ω.

`t` is a time instant

"""
function uncond_QFI_transverse(t, N, κ, ω)
    tot=0.
    for m=0:N
        r11= ρDiag(m,N,t,κ)
        r12=ρAntiDiag(m,N,t,κ, Complex(ω))
        dr12=dρAntiDiag(m,N,t,κ, Complex(ω))
        tot+= binomial(N,m)*QFIqubit(r11,r12,dr12)
    end
    return 0.5*real(tot)
end

ρDiag(m,N,t,κ) =
    0.5 .* (2.0^(-N) .* (1+-exp.(-κ .* t)) .^ (-m+N) .* (1.0+exp.(-κ .* t)) .^ m+2.0^(-N) .* (1-exp.(-κ .* t)) .^ m .* (1+exp.(-κ .* t)) .^ (-m+N))
ρAntiDiag(m,N,t,κ,ω) = (0.5) .* ((exp.((-0.5) .* κ .* t) .* κ .* (κ .^ 2-4.0 .* ω .^ 2) .^ (-0.5) .* sinh.((0.5) .* (κ .^ 2-4.0 .* ω .^ 2) .^ (0.5) .* t)) .^ m .* (exp.((-0.5) .* κ .* t) .* (cosh.((0.5) .* (κ .^ 2-4.0 .* ω .^ 2) .^ (0.5) .* t)+(1im*(-2)) .* ω .* (κ .^ 2-4.0 .* ω .^ 2) .^ (-0.5) .* sinh.((0.5) .* (κ .^ 2-4.0 .* ω .^ 2) .^ (0.5) .* t))) .^ (-m+N)+(exp.((-0.5) .* κ .* t) .* κ .* (κ .^ 2+(-4) .* ω .^ 2) .^ (-0.5) .* sinh.((0.5) .* (κ .^ 2-4.0 .* ω .^ 2) .^ (0.5) .* t)) .^ (-m+N) .* (exp.((-0.5) .* κ .* t) .* (cosh.((0.5) .* (κ .^ 2-4.0 .* ω .^ 2) .^ (0.5) .* t)+(1im*2) .* ω .* (κ .^ 2-4.0 .* ω .^ 2) .^ (-0.5) .* sinh.((0.5) .* (κ .^ 2-4.0 .* ω .^ 2) .^ (0.5) .* t))) .^ m)
dρAntiDiag(m,N,t,κ,ω)=-(κ .^ (2.0)+(-4) .* ω .^ (2.0)) .^ (-1.5) .* (exp.((-0.5) .* κ .* m .* t) .* m .* ω .* (κ .^ (2.0)+(-4) .* ω .^ (2.0)) .^ ((0.5)+(-0.5) .* m) .* ((-2)+(κ .^ (2.0)+(-4) .* ω .^ (2.0)) .^ (0.5) .* t .* coth.((0.5) .* (κ .^ (2.0)+(-4) .* ω .^ (2.0)) .^ (0.5) .* t)) .* (κ .* sinh.((0.5) .* (κ .^ (2.0)+(-4) .* ω .^ (2.0)) .^ (0.5) .* t)) .^ m .* (exp.((-0.5) .* κ .* t) .* (cosh.((0.5) .* (κ .^ (2.0)+(-4) .* ω .^ (2.0)) .^ (0.5) .* t)+(1im*(-2)) .* ω .* (κ .^ (2.0)+(-4) .* ω .^ (2.0)) .^ (-0.5) .* sinh.((0.5) .* (κ .^ (2.0)+(-4) .* ω .^ (2.0)) .^ (0.5) .* t))) .^ ((-1) .* m+N)+(-1) .* (m+(-1) .* N) .* ω .* (κ .^ (2.0)+(-4) .* ω .^ (2.0)) .^ (0.5) .* ((-2)+(κ .^ (2.0)+(-4) .* ω .^ (2.0)) .^ (0.5) .* t .* coth.((0.5) .* (κ .^ (2.0)+(-4) .* ω .^ (2.0)) .^ (0.5) .* t)) .* (exp.((-0.5) .* κ .* t) .* κ .* (κ .^ (2.0)+(-4) .* ω .^ (2.0)) .^ (-0.5) .* sinh.((0.5) .* (κ .^ (2.0)+(-4) .* ω .^ (2.0)) .^ (0.5) .* t)) .^ ((-1) .* m+N) .* (exp.((-0.5) .* κ .* t) .* (cosh.((0.5) .* (κ .^ (2.0)+(-4) .* ω .^ (2.0)) .^ (0.5) .* t)+(1im*2) .* ω .* (κ .^ (2.0)+(-4) .* ω .^ (2.0)) .^ (-0.5) .* sinh.((0.5) .* (κ .^ (2.0)+(-4) .* ω .^ (2.0)) .^ (0.5) .* t))) .^ m+1im .* exp.((-0.5) .* κ .* t) .* m .* (exp.((-0.5) .* κ .* t) .* κ .* (κ .^ (2.0)+(-4) .* ω .^ (2.0)) .^ (-0.5) .* sinh.((0.5) .* (κ .^ (2.0)+(-4) .* ω .^ (2.0)) .^ (0.5) .* t)) .^ ((-1) .* m+N) .* (exp.((-0.5) .* κ .* t) .* (cosh.((0.5) .* (κ .^ (2.0)+(-4) .* ω .^ (2.0)) .^ (0.5) .* t)+(1im*2) .* ω .* (κ .^ (2.0)+(-4) .* ω .^ (2.0)) .^ (-0.5) .* sinh.((0.5) .* (κ .^ (2.0)+(-4) .* ω .^ (2.0)) .^ (0.5) .* t))) .^ ((-1)+m) .* (2 .* ω .^ (2.0) .* (κ .^ (2.0)+(-4) .* ω .^ (2.0)) .^ (0.5) .* t .* cosh.((0.5) .* (κ .^ (2.0)+(-4) .* ω .^ (2.0)) .^ (0.5) .* t)+((1im*4) .* ω .^ 3 .* t+κ .^ (2.0) .* ((-1)+(1im*(-1)) .* ω .* t)) .* sinh.((0.5) .* (κ .^ (2.0)+(-4) .* ω .^ (2.0)) .^ (0.5) .* t))+1im .* exp.((-0.5) .* κ .* t) .* (m+(-1) .* N) .* (exp.((-0.5) .* κ .* t) .* κ .* (κ .^ (2.0)+(-4) .* ω .^ (2.0)) .^ (-0.5) .* sinh.((0.5) .* (κ .^ (2.0)+(-4) .* ω .^ (2.0)) .^ (0.5) .* t)) .^ m .* (exp.((-0.5) .* κ .* t) .* (cosh.((0.5) .* (κ .^ (2.0)+(-4) .* ω .^ (2.0)) .^ (0.5) .* t)+(1im*(-2)) .* ω .* (κ .^ (2.0)+(-4) .* ω .^ (2.0)) .^ (-0.5) .* sinh.((0.5) .* (κ .^ (2.0)+(-4) .* ω .^ (2.0)) .^ (0.5) .* t))) .^ ((-1)+(-1) .* m+N) .* (2 .* ω .^ (2.0) .* (κ .^ (2.0)+(-4) .* ω .^ (2.0)) .^ (0.5) .* t .* cosh.((0.5) .* (κ .^ (2.0)+(-4) .* ω .^ (2.0)) .^ (0.5) .* t)+1im .* ((-4) .* ω .^ 3 .* t+κ .^ (2.0) .* (1im+ω .* t)) .* sinh.((0.5) .* (κ .^ (2.0)+(-4) .* ω .^ (2.0)) .^ (0.5) .* t)))

function QFIqubit(ρ11,ρ12,dρ12)
    reC=real(ρ12./ρ11)
    imC=imag(ρ12./ρ11)
    redC=real(dρ12./ρ11)
    imdC=imag(dρ12./ρ11)
    return 2*ρ11 .* (redC .^ 2 + imdC .^ 2 + ( (reC .* redC + imC .* imdC) .^ 2)./(1 - reC .^ 2 - imC .^ 2))
end


text:
module Discretization

include("tensor_product.jl")

end


text:
function EnclosingBall(arg0::Point, arg1::jdouble, arg2::Vector{Point})
    return EnclosingBall((Point, jdouble, Vector{Point}), arg0, arg1, arg2)
end

function contains(obj::EnclosingBall, arg0::Point)
    return jcall(obj, "contains", jboolean, (Point,), arg0)
end

function contains(obj::EnclosingBall, arg0::Point, arg1::jdouble)
    return jcall(obj, "contains", jboolean, (Point, jdouble), arg0, arg1)
end

function equals(obj::Object, arg0::Object)
    return jcall(obj, "equals", jboolean, (Object,), arg0)
end

function get_center(obj::EnclosingBall)
    return jcall(obj, "getCenter", Point, ())
end

function get_class(obj::Object)
    return jcall(obj, "getClass", Class, ())
end

function get_radius(obj::EnclosingBall)
    return jcall(obj, "getRadius", jdouble, ())
end

function get_support(obj::EnclosingBall)
    return jcall(obj, "getSupport", Vector{Point}, ())
end

function get_support_size(obj::EnclosingBall)
    return jcall(obj, "getSupportSize", jint, ())
end

function hash_code(obj::Object)
    return jcall(obj, "hashCode", jint, ())
end

function notify(obj::Object)
    return jcall(obj, "notify", void, ())
end

function notify_all(obj::Object)
    return jcall(obj, "notifyAll", void, ())
end

function to_string(obj::Object)
    return jcall(obj, "toString", JString, ())
end

function wait(obj::Object)
    return jcall(obj, "wait", void, ())
end

function wait(obj::Object, arg0::jlong)
    return jcall(obj, "wait", void, (jlong,), arg0)
end

function wait(obj::Object, arg0::jlong, arg1::jint)
    return jcall(obj, "wait", void, (jlong, jint), arg0, arg1)
end


text:
using Test
import AMRVW
const A = AMRVW
using LinearAlgebra

@testset "passthrough!-chains"  begin

    T = Float64
    S = Complex{T}

    # check passthorugh! with Desc; Asc; and Twisted
    N = 10
    Des = A.random_rotator.(S, 1:N) |> A.DescendingChain
    Asc = A.random_rotator.(S, N:-1:1) |> A.AscendingChain
    #Tw = A.random_rotator.(S, 0 .+ reverse([1,3,2,4,7,6,5 ])) |> A.TwistedChain
    Tw = A.random_rotator.(S, 0 .+ reverse([3,4,7,6,5 ])) |> A.TwistedChain

    M = diagm(0 => ones(S, N+1))

    ## We have Ascending, Descending, and Twisted to check
    N = 10
    Des = A.random_rotator.(S, 1:N) |> A.DescendingChain
    Asc = A.random_rotator.(S, N-1:-1:1) |> A.AscendingChain
    Tw = A.random_rotator.(S, 0 .+ reverse([3,4,7,6,5 ])) |> A.TwistedChain

    # <--
    M1 = Des * (Tw * M)
    A.passthrough!(Des, Tw)
    M2 = Tw * (Des * M)
    @test M1 - M2 .|> round5 .|> iszero |> all

    M1 = Asc * (Tw * M)
    A.passthrough!(Asc, Tw)
    M2 = Tw * (Asc * M)
    @test M1 - M2 .|> round5 .|> iszero |> all

    # -->
    M1 = Tw * (Des * M)
    A.passthrough!(Tw, Des)
    M2 = Des * (Tw * M)
    @test M1 - M2 .|> round5 .|> iszero |> all

    M1 = Tw * (Asc * M)
    A.passthrough!(Tw, Asc)
    M2 = Asc * (Tw * M)
    @test M1 - M2 .|> round5 .|> iszero |> all

    ## check passthorugh Des, Asc
    M1 = Des * (Asc * M)
    A.passthrough!(Des, Asc)
    M2 = Asc * (Des * M)
    @test M1 - M2 .|> round5 .|> iszero |> all

    Des = A.random_rotator.(S, 1:N-1) |> A.DescendingChain
    Asc = A.random_rotator.(S, N:-1:1) |> A.AscendingChain
    M1 = Des * (Asc * M)
    A.passthrough!(Des, Asc)
    M2 = Asc * (Des * M)
    @test M1 - M2 .|> round5 .|> iszero |> all

    Des = A.random_rotator.(S, 1:N) |> A.DescendingChain
    Asc = A.random_rotator.(S, N:-1:2) |> A.AscendingChain
    M1 = Asc * (Des * M)
    A.passthrough!(Asc, Des)
    M2 = Des * (Asc * M)
    @test M1 - M2 .|> round5 .|> iszero |> all

    Des = A.random_rotator.(S, 2:N) |> A.DescendingChain
    Asc = A.random_rotator.(S, N:-1:1) |> A.AscendingChain
    M1 = Asc * (Des * M)
    A.passthrough!(Asc, Des)
    M2 = Des * (Asc * M)
    @test M1 - M2 .|> round5 .|> iszero |> all
end

@testset "passthrough_phase!"  begin

    # test passthrough_phase!
    T = Float64; S = Complex{T};
    n = 4

    ## Descending
    for i in 1:6
        Ms = A.random_rotator.(S,  [2,3,4])
        Des = A.random_rotator.(S,  [1,2,3,4,5])
        D =  A.SparseDiagonal(S,6+1)
        Di = A.DiagonalRotator(complex(sincos(rand())...), i)
        M1 = Di * (Ms * (Des *diagm(0 => D.x)))

        A.passthrough_phase!(Di, A.DescendingChain(Ms), (A.DescendingChain(Des),), D)

        M2 = Ms * (Des *diagm(0 => D.x))
        @test M1 - M2 |> round5  .|> iszero |> all
    end

    ## Ascending
    for i in 1:6
        Ms = A.random_rotator.(S,  [4,3,2])
        Des = A.random_rotator.(S,  [1,2,3,4,5])
        D =  A.SparseDiagonal(S,6+1)
        Di = A.DiagonalRotator(complex(sincos(rand())...), i)
        M1 = Di * (Ms * (Des *diagm(0 => D.x)))

        A.passthrough_phase!(Di, A.AscendingChain(Ms), (A.DescendingChain(Des),), D)

        M2 = Ms * (Des *diagm(0 => D.x))
        @test M1 - M2 |> round5  .|> iszero |> all
    end

    ## Twisted
    for i in 1:6
        Ms = A.TwistedChain(A.random_rotator.(S,  [2,3,4,5]))
        D =  A.SparseDiagonal(S,6+1)
        Di = A.DiagonalRotator(complex(sincos(rand())...), i)
        M1 = Di * (Ms * diagm(0 => D.x))

        A.passthrough_phase!(Di, Ms, (), D)

        M2 = Ms * diagm(0 => D.x)
        @test M1 - M2 |> round5  .|> iszero |> all
    end

    ##
    for i in 1:6
        Ms = A.TwistedChain(A.random_rotator.(S,  [5,4,3,2]))
        D =  A.SparseDiagonal(S,6+1)
        Di = A.DiagonalRotator(complex(sincos(rand())...), i)
        M1 = Di * (Ms * diagm(0 => D.x))

        A.passthrough_phase!(Di, Ms, (), D)

        M2 = Ms * diagm(0 => D.x)
        @test M1 - M2 |> round5  .|> iszero |> all
    end

    ## CMV
    for i in 1:6
        Ms = A.TwistedChain(A.random_rotator.(S,  [2,4,3,5])) # XX
        Des = A.random_rotator.(S,  [1,2,3,4,5])
        D =  A.SparseDiagonal(S,6+1)
        Di = A.DiagonalRotator(complex(sincos(rand())...), i)
        M1 = Di * (Ms * (Des *diagm(0 => D.x)))

        A.passthrough_phase!(Di, Ms, (A.DescendingChain(Des),), D)

        M2 = Ms * (Des * diagm(0 => D.x))
        @test M1 - M2 |> round5  .|> iszero |> all
    end

    ##
    for i in 1:6

        Ms = A.TwistedChain(A.random_rotator.(S,  [3,2,4,5])) # XXX 5 is issue
        Des = A.random_rotator.(S,  [1,2,3,4,5])
        D =  A.SparseDiagonal(S,6+1)
        Di = A.DiagonalRotator(complex(sincos(rand())...), i)
        M1 = Di * (Ms * (Des * diagm(0 => D.x)))

        A.passthrough_phase!(Di, Ms, (A.DescendingChain(Des),), D)

        M2 = Ms * (Des * diagm(0 => D.x))
        @test M1 - M2 |> round5  .|> iszero |> all
    end

end


text:
module Sounds

using Unitful
using Requires

include(joinpath(@__DIR__,"units.jl"))
include(joinpath(@__DIR__,"sound.jl"))
include(joinpath(@__DIR__,"samplerate.jl"))
include(joinpath(@__DIR__,"audible.jl"))

const localunits = Unitful.basefactors
const localpromotion = Unitful.promotion
function __init__()
  merge!(Unitful.basefactors,localunits)
  merge!(Unitful.promotion, localpromotion)
end

@require PortAudio begin
  import PortAudio
  import SampledSignals
  export play

  const portaudio = PortAudio.PortAudioStream()
  Base.write(io::PortAudio.PortAudioStream,x::Sound) =
    write(io,SampledSignals.SampleBuf(x))
  play(x::Sound) = write(portaudio,x)
  atexit() do
    sleep(0.1)
    close(portaudio)
  end
end

end


text:
function LoadTestToyData(Oracles)

	features =
	[-1.0  1.0
	-1.0  -1.0
	-1.0  -2.0
	-1.0  -3.0
	1.0  -1.0
	-2.0  -2.0
	1.0   1.0
	-2.0  -3.0
	-0.9  -0.9
	-1.1  -2.1
	-0.8  -0.8
	-1.2  -2.2
	-0.7  -0.7
	-1.3  -2.3
	-0.6  -0.6
	-1.4  -2.4]

	labels = [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.0,-1.0,1.0,-1.0]

	normalizeFeatures(features, size(features)[2])

	(classLabels, numClasses) = createClassLabels(labels)

	(trf, trl, trl2, numTP, tef, tel, tel2) = trainTestRandomSeparate(features, MinusPlusOneVector(labels, Set([1.0])), classLabels)

	L2regs=[
	false,
	true
	]
	for L2reg in L2regs

		(gradientOracle, numVars, numTP, csDataType, LossFunctionString, myOutputter, L2reg) = createBLOracles(trf, trl, numTP, tef, tel, L2reg; outputLevel = 0)

		push!(Oracles, (gradientOracle, numVars, numTP, csDataType, LossFunctionString, myOutputter, L2reg,  "TestToy", (t)-> Problem(L2reg, "TestToy", LossFunctionString, numTP, t,(j)->getSampleFunctionAt(j,gradientOracle))))

		(gradientOracle2, numVars2, numTP2, csDataType2, LossFunctionString2, myOutputter2, L2reg2) = createMLOracles(trf, trl2, numTP, numClasses, tef, tel2, L2reg; outputLevel = 0)

		push!(Oracles,(gradientOracle2, numVars2, numTP2, csDataType2, LossFunctionString2, myOutputter2, L2reg2,  "TestToy",  (t)-> Problem(L2reg2, "TestToy", LossFunctionString2,  numTP2, t ,(j)->getSampleFunctionAt(j,gradientOracle2))))

	end
end

LoadTestToyData(Oracles)


text:
"""
Series - ?

TODO add function signatures
TODO add description

## Attributes
$(ATTRIBUTES)
"""
@recipe(Series, series) do scene
Attributes(
    seriescolors = :Set1,
    seriestype = :lines
)
end

convert_arguments(::Type{<: Series}, A::AbstractMatrix{<: Number}) = (A,)

function plot!(sub::Series)
    A = sub[1]
    colors = map_once(sub[:seriescolors], A) do colors, A
        cmap = to_colormap(colors)
        if size(A, 2) > length(cmap)
            @info("Colormap doesn't have enough distinctive values. Please consider using another value for seriescolors")
            cmap = interpolated_getindex.((cmap,), range(0, stop=1, length=M))
        end
        cmap
    end
    map_once(A, sub[:seriestype]) do A, stype
        empty!(sub.plots)
        N, M = size(A)
        for i = 1:M
            c = lift(getindex, colors, i)
            attributes = Attributes(color = c)
            a_view = view(A, :, i)
            if stype in (:lines, :scatter_lines)
                lines!(sub, attributes, a_view)
            end
            # if stype in (:scatter, :scatter_lines)
            #     scatter!(subsub, attributes, a_view)
            # end
            # subsub
        end
    end
    labels = get(sub, :labels) do
        map(i-> "y $i", 1:size(A[], 2))
    end
    oldlegend!(sub, copy(sub.plots), labels)
    sub
end


text:
using Test

using MarkovChains
@testset "poisson_trunc" begin
    λ = 20.0
    tol = 1e-6
    ltp, rtp, probs = poisson_trunc_point(λ, tol)
    @test ltp == 3
    @test rtp == 45
    @test abs(1.0 - sum(probs)) < tol
    λ = 400.0
    ltp, rtp, probs = poisson_trunc_point(λ, tol)
    @test ltp == 306
    @test rtp == 502
    @test abs(1.0 - sum(probs)) < tol
end


text:
@recipe function f(sresults::Simulation)
    DF = sresults.sampleddata.DF
    VAF = DF[:VAF]
    freq = DF[:freq]

    if length(sresults.output.clonefreq) > 0
        xint = sresults.output.clonefreq./2 * sresults.input.cellularity

        @series begin
            seriestype --> :bar
            linecolor --> :darkslategrey
            fillcolor --> :darkslategrey
            markerstrokecolor --> :white
            legend --> false
            grid --> false
            VAF, freq
        end

        @series begin
            yaxis --> "Counts"
            xaxis --> "VAF"
            seriestype --> :vline
            legend --> false
            fillcolor --> :darkred
            linewidth --> 3
            xint
        end
    else
        @series begin
            seriestype --> :bar
            yaxis --> "Counts"
            xaxis --> "VAF"
            linecolor --> :darkslategrey
            fillcolor --> :darkslategrey
            markerstrokecolor --> :white
            legend --> false
            grid --> false
            VAF, freq
        end
    end
end


text:
Base.@kwdef mutable struct Path
    directory::String = ""
    filename::String = ""
end

function get_directory(p::Path)
    p.directory
end

function get_filename(p::Path)
    p.filename
end

function set_directory!(p::Path, directory::String)
    p.directory = directory
end

function set_filename!(p::Path, filename::String)
    p.filename = filename
end


text:
__precompile__()

module CpelAsm
###################################################################################################
# DEPENDENCIES
###################################################################################################
using Documenter                # For documentation
using Distributed               # For parallelization
using SharedArrays              # For parallelization
using DelimitedFiles            # For delimited files
using Random                    # For random number generation
using StatsBase                 # For statistics
using Calculus                  # For Hessian estimation
using LinearAlgebra             # For Eigendecomposition of Σ
using Dates                     # For showing time
using Distributions             # For sampling RVs
using LaTeXStrings              # For labels in plots
using ForwardDiff               # For auto-differentiation
using Optim                     # For Simulated annealing
using NLsolve                   # For EM algorithm
using BioAlignments             # For reading-in BAM/SAM files
using GenomicFeatures           # For BED, GFF3, and bigWig files
using GeneticVariation          # For VCF file
using BioSequences              # For FASTA file
using Combinatorics             # For permutation test
using MultipleTesting           # For multiple hypothesis testing
###################################################################################################
# INCLUDES
###################################################################################################
include("CpelAsmEstimation.jl")
include("CpelAsmSimulations.jl")
include("CpelAsmBioinformatics.jl")
# include("CpelAsmAgnostic.jl")
include("CpelAsmOutput.jl")
###################################################################################################
# EXPORTS
###################################################################################################

# Data generating functions
export gen_x_mc                     # Generate a methylation vector using MC approach
export gen_ising_full_data          # Generate a set of full observations using slow approach
export gen_ising_part_data          # Generate a set of partial observations using slow approach

# Different computations
export comp_Z                       # Partition function
export est_alpha                    # Estimate α for N=1 case
export est_theta_sa                 # Estimate θ in general case using Simulated Annealing
export est_theta_em                 # Estimate θ in general case using EM algorithm
export comp_lkhd                    # Compute likelihood
export comp_ex                      # Compute E[X] vector
export comp_exx                     # Compute E[XX] vector
export comp_cov                     # Compute Σ matrix
export comp_mml                     # Compute MML in `comp_tobs`
export comp_mml_∇                   # Compute MML in `comp_tnull (uses ∇logZ)
export comp_nme                     # Compute NME in `comp_tobs`
export comp_nme_∇                   # Compute NME in `comp_tnull` (uses ∇logZ)
export comp_nme_mix_mc              # Compute NME in mixture model using MC technique
export comp_nme_mix_exact           # Compute NME in mixture model recursively (N<17)
export comp_uc                      # Compute UC

# General functions
export gen_gffs                     # Function to generate GFF files
export comp_tobs                    # Compute statistis in haplotypes
export comp_tnull                   # Generate null statistis
export comp_pvals                   # Compute p-values
export run_analysis                 # Run it all
# export run_uc_analysis              # Run allele-agnostic UC computations
# export run_jsd_analysis             # Run allele-agnostic JSD computations
# export run_allele_agnostic_analysis # Run allele-agnostic analysis

end # module


text:
# ---
# title: 375. Guess Number Higher or Lower II
# id: problem375
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Dynamic Programming, Minimax
# link: <https://leetcode.com/problems/guess-number-higher-or-lower-ii/description/>
# hidden: true
# ---
#
# We are playing the Guessing Game. The game will work as follows:
#
#   1. I pick a number between `1` and `n`.
#   2. You guess a number.
#   3. If you guess the right number, **you win the game**.
#   4. If you guess the wrong number, then I will tell you whether the number I picked is **higher or lower** , and you will continue guessing.
#   5. Every time you guess a wrong number `x`, you will pay `x` dollars. If you run out of money, **you lose the game**.
#
# Given a particular `n`, return  _the minimum amount of money you need to
# **guarantee a win regardless of what number I pick**_.
#
#
#
# **Example 1:**
#
# ![](https://assets.leetcode.com/uploads/2020/09/10/graph.png)
#
#
#
#     Input: n = 10
#     Output: 16
#     Explanation: The winning strategy is as follows:
#     - The range is [1,10]. Guess 7.
#         - If this is my number, your total is /$/0. Otherwise, you pay /$/7.
#         - If my number is higher, the range is [8,10]. Guess 9.
#             - If this is my number, your total is /$/7. Otherwise, you pay /$/9.
#             - If my number is higher, it must be 10. Guess 10. Your total is /$/7 + /$/9 = /$/16.
#             - If my number is lower, it must be 8. Guess 8. Your total is /$/7 + /$/9 = /$/16.
#         - If my number is lower, the range is [1,6]. Guess 3.
#             - If this is my number, your total is /$/7. Otherwise, you pay /$/3.
#             - If my number is higher, the range is [4,6]. Guess 5.
#                 - If this is my number, your total is /$/7 + /$/3 = /$/10. Otherwise, you pay /$/5.
#                 - If my number is higher, it must be 6. Guess 6. Your total is /$/7 + /$/3 + /$/5 = /$/15.
#                 - If my number is lower, it must be 4. Guess 4. Your total is /$/7 + /$/3 + /$/5 = /$/15.
#             - If my number is lower, the range is [1,2]. Guess 1.
#                 - If this is my number, your total is /$/7 + /$/3 = /$/10. Otherwise, you pay /$/1.
#                 - If my number is higher, it must be 2. Guess 2. Your total is /$/7 + /$/3 + /$/1 = /$/11.
#     The worst case in all these scenarios is that you pay /$/16. Hence, you only need /$/16 to guarantee a win.
#
#
# **Example 2:**
#
#
#
#     Input: n = 1
#     Output: 0
#     Explanation:  There is only one possible number, so you can guess 1 and not have to pay anything.
#
#
# **Example 3:**
#
#
#
#     Input: n = 2
#     Output: 1
#     Explanation:  There are two possible numbers, 1 and 2.
#     - Guess 1.
#         - If this is my number, your total is /$/0. Otherwise, you pay /$/1.
#         - If my number is higher, it must be 2. Guess 2. Your total is /$/1.
#     The worst case is that you pay /$/1.
#
#
#
#
# **Constraints:**
#
#   * `1 <= n <= 200`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


text:
utc = FixedTimeZone("UTC")
fixed = FixedTimeZone("UTC-06:00")
winnipeg = first(compile("America/Winnipeg", tzdata["northamerica"]))   # UTC-6:00 (or UTC-5:00)
st_johns = first(compile("America/St_Johns", tzdata["northamerica"]))   # UTC-3:30 (or UTC-2:30)
eucla = first(compile("Australia/Eucla", tzdata["australasia"]))        # UTC+8:45
colombo = first(compile("Asia/Colombo", tzdata["asia"]))                # See note below

# On 1996-05-25 at 00:00, the Asia/Colombo time zone in Sri Lanka moved from Indian Standard
# Time (UTC+5:30) to Lanka Time (UTC+6:30). On 1996-10-26 at 00:30, Lanka Time was revised
# from UTC+6:30 to UTC+6:00, marking a -00:30 transition. Transitions like these are doubly
# unusual (compared to the more common DST transitions) as it is both an half-hour
# transition and a transition that lands at midnight (causing 1996-10-26T00:00 to be
# ambiguous; midnights are rarely ambiguous). In 2006, Asia/Colombo returned to Indian
# Standard Time, causing another -00:30 transition from 00:30 to 00:00.

##################
# NO TRANSITIONS #
##################

# Test rounding where no rounding is necessary.

dt = DateTime(2016)

for tz in [utc, fixed, winnipeg, st_johns, eucla, colombo]
    local tz
    local zdt = ZonedDateTime(dt, tz)
    for p in [Dates.Year, Dates.Month, Dates.Day, Dates.Hour, Dates.Minute, Dates.Second]
        @test floor(zdt, p) == zdt
        @test ceil(zdt, p) == zdt
        @test round(zdt, p) == zdt
    end
end

# Test rounding non-controversial ZonedDateTimes (no transitions).

dt = DateTime(2016, 2, 5, 13, 10, 20, 500)

for tz in [utc, fixed, winnipeg, st_johns, eucla, colombo]
    local tz
    local zdt = ZonedDateTime(dt, tz)

    @test floor(zdt, Dates.Year) == ZonedDateTime(2016, tz)
    @test floor(zdt, Dates.Month) == ZonedDateTime(2016, 2, tz)
    @test floor(zdt, Dates.Week) == ZonedDateTime(2016, 2, tz)      # Previous Monday
    @test floor(zdt, Dates.Day) == ZonedDateTime(2016, 2, 5, tz)
    @test floor(zdt, Dates.Hour) == ZonedDateTime(2016, 2, 5, 13, tz)
    @test floor(zdt, Dates.Minute) == ZonedDateTime(2016, 2, 5, 13, 10, tz)
    @test floor(zdt, Dates.Second) == ZonedDateTime(2016, 2, 5, 13, 10, 20, tz)

    @test ceil(zdt, Dates.Year) == ZonedDateTime(2017, tz)
    @test ceil(zdt, Dates.Month) == ZonedDateTime(2016, 3, tz)
    @test ceil(zdt, Dates.Week) == ZonedDateTime(2016, 2, 8, tz)    # Following Monday
    @test ceil(zdt, Dates.Day) == ZonedDateTime(2016, 2, 6, tz)
    @test ceil(zdt, Dates.Hour) == ZonedDateTime(2016, 2, 5, 14, tz)
    @test ceil(zdt, Dates.Minute) == ZonedDateTime(2016, 2, 5, 13, 11, tz)
    @test ceil(zdt, Dates.Second) == ZonedDateTime(2016, 2, 5, 13, 10, 21, tz)

    @test round(zdt, Dates.Year) == ZonedDateTime(2016, tz)
    @test round(zdt, Dates.Month) == ZonedDateTime(2016, 2, tz)
    @test round(zdt, Dates.Week) == ZonedDateTime(2016, 2, 8, tz)   # Following Monday
    @test round(zdt, Dates.Day) == ZonedDateTime(2016, 2, 6, tz)
    @test round(zdt, Dates.Hour) == ZonedDateTime(2016, 2, 5, 13, tz)
    @test round(zdt, Dates.Minute) == ZonedDateTime(2016, 2, 5, 13, 10, tz)
    @test round(zdt, Dates.Second) == ZonedDateTime(2016, 2, 5, 13, 10, 21, tz)
end

##########################
# DST TRANSITION FORWARD #
##########################

# Test rounding over spring transition (missing hour). FixedTimeZones have no transitions,
# but ZonedDateTimes with VariableTimeZones will round in their current (fixed) time zone
# and then adjust to the new time zone if a transition has occurred (DST, for example).

# Test rounding backward, toward the missing hour.

dt = DateTime(2016, 3, 13, 3, 15)             # 15 minutes after transition

zdt = ZonedDateTime(dt, fixed)
@test floor(zdt, Dates.Day) == ZonedDateTime(2016, 3, 13, fixed)
@test floor(zdt, Dates.Hour(2)) == ZonedDateTime(2016, 3, 13, 2, fixed)

for tz in [winnipeg, st_johns]
    local tz
    local zdt = ZonedDateTime(dt, tz)
    @test floor(zdt, Dates.Day) == ZonedDateTime(2016, 3, 13, tz)
    @test floor(zdt, Dates.Hour(2)) == ZonedDateTime(2016, 3, 13, 1, tz)
end

# Test rounding forward, toward the missing hour.

dt = DateTime(2016, 3, 13, 1, 55)             # 5 minutes before transition

zdt = ZonedDateTime(dt, fixed)
@test ceil(zdt, Dates.Day) == ZonedDateTime(2016, 3, 14, fixed)
@test ceil(zdt, Dates.Hour) == ZonedDateTime(2016, 3, 13, 2, fixed)
@test ceil(zdt, Dates.Minute(30)) == ZonedDateTime(2016, 3, 13, 2, fixed)
@test round(zdt, Dates.Minute(30)) == ZonedDateTime(2016, 3, 13, 2, fixed)

for tz in [winnipeg, st_johns]
    local tz
    local zdt = ZonedDateTime(dt, tz)

    @test ceil(zdt, Dates.Day) == ZonedDateTime(2016, 3, 14, tz)
    @test ceil(zdt, Dates.Hour) == ZonedDateTime(2016, 3, 13, 3, tz)
    @test ceil(zdt, Dates.Minute(30)) == ZonedDateTime(2016, 3, 13, 3, tz)
    @test round(zdt, Dates.Minute(30)) == ZonedDateTime(2016, 3, 13, 3, tz)
end

# Test rounding from "midday".

dt = DateTime(2016, 3, 13, 12)                # Noon on day of transition

# Noon is the middle of the day, and ties round up by default.
@test round(ZonedDateTime(dt, fixed), Dates.Day) == ZonedDateTime(2016, 3, 14, fixed)

# Noon isn't the middle of the day, as 2:00 through 2:59:59.999 are missing in these zones.
@test round(ZonedDateTime(dt, winnipeg), Dates.Day) == ZonedDateTime(2016, 3, 13, winnipeg)
@test round(ZonedDateTime(dt, st_johns), Dates.Day) == ZonedDateTime(2016, 3, 13, st_johns)

###########################
# DST TRANSITION BACKWARD #
###########################

# Test rounding over autumn transition (additional, ambiguous hour).

# Test rounding backward, toward the ambiguous hour.

dt = DateTime(2015, 11, 1, 2, 15)             # 15 minutes after ambiguous hour

zdt = ZonedDateTime(dt, fixed)
@test floor(zdt, Dates.Day) == ZonedDateTime(2015, 11, 1, fixed)
@test floor(zdt, Dates.Hour(3)) == ZonedDateTime(2015, 11, 1, fixed)
# Rounding to Hour(3) will give 00:00, 03:00, 06:00, 09:00, etc.

for tz in [winnipeg, st_johns]
    local tz
    local zdt = ZonedDateTime(dt, tz)
    @test floor(zdt, Dates.Day) == ZonedDateTime(2015, 11, 1, tz)
    @test floor(zdt, Dates.Hour(3)) == ZonedDateTime(2015, 11, 1, 1, tz, 1)
    # Rounding is performed in the current fixed zone, then relocalized if a transition has
    # occurred. This means that instead of 00:00, 03:00, etc., we expect 01:00, 04:00, etc.
end

# Test rounding forward, toward the ambiguous hour.

dt = DateTime(2015, 11, 1, 0, 55)             # 5 minutes before ambiguous hour

zdt = ZonedDateTime(dt, fixed)
@test ceil(zdt, Dates.Day) == ZonedDateTime(2015, 11, 2, fixed)
@test ceil(zdt, Dates.Hour) == ZonedDateTime(2015, 11, 1, 1, fixed)
@test ceil(zdt, Dates.Minute(30)) == ZonedDateTime(2015, 11, 1, 1, fixed)
@test round(zdt, Dates.Minute(30)) == ZonedDateTime(2015, 11, 1, 1, fixed)

for tz in [winnipeg, st_johns]
    local tz
    local zdt = ZonedDateTime(dt, tz)
    next_hour = ZonedDateTime(DateTime(2015, 11, 1, 1), tz, 1)

    @test ceil(zdt, Dates.Day) == ZonedDateTime(2015, 11, 2, tz)
    @test ceil(zdt, Dates.Hour) == next_hour
    @test ceil(zdt, Dates.Minute(30)) == next_hour
    @test round(zdt, Dates.Minute(30)) == next_hour
end

# Test rounding forward and backward, during the ambiguous hour.

dt = DateTime(2015, 11, 1, 1, 25)                   # During ambiguous hour

zdt = ZonedDateTime(dt, fixed)
@test floor(zdt, Dates.Day) == ZonedDateTime(2015, 11, 1, fixed)
@test ceil(zdt, Dates.Day) == ZonedDateTime(2015, 11, 2, fixed)
@test round(zdt, Dates.Day) == ZonedDateTime(2015, 11, 1, fixed)
@test floor(zdt, Dates.Hour) == ZonedDateTime(2015, 11, 1, 1, fixed)
@test ceil(zdt, Dates.Hour) == ZonedDateTime(2015, 11, 1, 2, fixed)
@test round(zdt, Dates.Hour) == ZonedDateTime(2015, 11, 1, 1, fixed)
@test floor(zdt, Dates.Minute(30)) == ZonedDateTime(2015, 11, 1, 1, fixed)
@test ceil(zdt, Dates.Minute(30)) == ZonedDateTime(2015, 11, 1, 1, 30, fixed)
@test round(zdt, Dates.Minute(30)) == ZonedDateTime(2015, 11, 1, 1, 30, fixed)

for tz in [winnipeg, st_johns]
    local tz
    local zdt = ZonedDateTime(dt, tz, 1)            # First 1:25, before "falling back"
    prev_hour = ZonedDateTime(2015, 11, 1, 1, tz, 1)
    between_hours = ZonedDateTime(2015, 11, 1, 1, 30, tz, 1)
    next_hour = ZonedDateTime(2015, 11, 1, 1, tz, 2)
    @test floor(zdt, Dates.Day) == ZonedDateTime(2015, 11, 1, tz)
    @test ceil(zdt, Dates.Day) == ZonedDateTime(2015, 11, 2, tz)
    @test round(zdt, Dates.Day) == ZonedDateTime(2015, 11, 1, tz)
    @test floor(zdt, Dates.Hour) == prev_hour
    @test ceil(zdt, Dates.Hour) == next_hour
    @test round(zdt, Dates.Hour) == prev_hour
    @test floor(zdt, Dates.Minute(30)) == prev_hour
    @test ceil(zdt, Dates.Minute(30)) == between_hours
    @test round(zdt, Dates.Minute(30)) == between_hours
end

###########################
# ASIA/COLOMBO TRANSITION #
###########################

# Test rounding to ambiguous midnight, which (unfortunately) isn't handled well when
# rounding to a DatePeriod resolution.

zdt = ZonedDateTime(1996, 10, 25, 23, 55, colombo)  # 5 minutes before ambiguous half-hour
@test floor(zdt, Dates.Day) == ZonedDateTime(1996, 10, 25, colombo)
@test_throws AmbiguousTimeError ceil(zdt, Dates.Day)
@test_throws AmbiguousTimeError round(zdt, Dates.Day)

zdt = ZonedDateTime(1996, 10, 26, 0, 35, colombo)   # 5 minutes after ambiguous half-hour
@test_throws AmbiguousTimeError floor(zdt, Dates.Day)
@test ceil(zdt, Dates.Day) == ZonedDateTime(1996, 10, 27, colombo)
@test_throws AmbiguousTimeError round(zdt, Dates.Day)

# Rounding to the ambiguous midnight works fine using a TimePeriod resolution, however.

zdt = ZonedDateTime(1996, 10, 25, 23, 55, colombo)  # 5 minutes before ambiguous half-hour
@test ceil(zdt, Dates.Hour) == ZonedDateTime(1996, 10, 26, colombo, 1)
@test round(zdt, Dates.Hour) == ZonedDateTime(1996, 10, 26, colombo, 1)

zdt = ZonedDateTime(1996, 10, 26, 0, 35, colombo)   # 5 minutes after ambiguous half-hour
@test floor(zdt, Dates.Hour) == ZonedDateTime(1996, 10, 26, colombo, 2)

# Rounding during the first half-hour between 00:00 and 00:30.

zdt = ZonedDateTime(1996, 10, 26, 0, 15, colombo, 1)
@test floor(zdt, Dates.Hour) == ZonedDateTime(1996, 10, 26, colombo, 1)
@test ceil(zdt, Dates.Hour) == ZonedDateTime(1996, 10, 26, 0, 30, colombo)
@test round(zdt, Dates.Hour) == ZonedDateTime(1996, 10, 26, colombo, 1)
@test floor(zdt, Dates.Minute(30)) == ZonedDateTime(1996, 10, 26, colombo, 1)
@test ceil(zdt, Dates.Minute(30)) == ZonedDateTime(1996, 10, 26, colombo, 2)
@test round(zdt, Dates.Minute(30)) == ZonedDateTime(1996, 10, 26, colombo, 2)

# Rounding during the second half-hour between 00:00 and 00:30.

zdt = ZonedDateTime(1996, 10, 26, 0, 15, colombo, 2)
@test floor(zdt, Dates.Hour) == ZonedDateTime(1996, 10, 26, colombo, 2)
@test ceil(zdt, Dates.Hour) == ZonedDateTime(1996, 10, 26, 1, colombo)
@test round(zdt, Dates.Hour) == ZonedDateTime(1996, 10, 26, colombo, 2)
@test floor(zdt, Dates.Minute(30)) == ZonedDateTime(1996, 10, 26, colombo, 2)
@test ceil(zdt, Dates.Minute(30)) == ZonedDateTime(1996, 10, 26, 0, 30, colombo)
@test round(zdt, Dates.Minute(30)) == ZonedDateTime(1996, 10, 26, 0, 30, colombo)

###############
# ERROR CASES #
###############

# Test rounding to invalid resolutions.

dt = DateTime(2016, 2, 28, 12, 15, 10, 190)
for tz in [utc, fixed, winnipeg, st_johns, eucla, colombo]
    local tz
    local zdt = ZonedDateTime(dt, tz)
    for p in [Dates.Year, Dates.Month, Dates.Day, Dates.Hour, Dates.Minute, Dates.Second]
        for v in [-1, 0]
            @test_throws DomainError floor(dt, p(v))
            @test_throws DomainError ceil(dt, p(v))
            @test_throws DomainError round(dt, p(v))
        end
    end
end


text:
@testset "Fonts" begin
    @testset "Caching" begin
        for font in values(_default_fonts)
            @test load_font(font) === load_font(font)
        end
    end
end