include("multigrid/utils.jl")
include("multigrid/poisson.jl")
include("multigrid/transfer.jl")
include("multigrid/smoothers.jl")
include("multigrid/multigrid.jl")
include("mgpcg.jl")

function solve_poisson_mgpcg(
    d2u::AbstractArray{<:AbstractFloat, N},
    mask::AbstractArray{Bool, 3},
    dx::NTuple{3, Real};
    kwargs...
) where {N}
    N ∈ (3, 4) || throw(ArgumentError("arrays must be 3d or 4d, got $(N)d"))
    return solve_poisson_mgpcg!(tzero(d2u), d2u, mask, dx; kwargs...)
end

function solve_poisson_mgpcg!(
    u::AbstractArray{<:AbstractFloat, N},
    d2u::AbstractArray{T, N},
    mask::AbstractArray{Bool, 3},
    dx::NTuple{3, Real};
    kwargs...
) where {T, N}
    N ∈ (3, 4) || throw(ArgumentError("arrays must be 3d or 4d, got $(N)d"))
    size(u) == size(d2u) || throw(DimensionMismatch())
    size(mask) == size(u)[1:3] || throw(DimensionMismatch())

    # crop to avoid unnecessary work. u[!mask] = 0
    sz = size(mask)
    Rc = crop_indices(mask)

    # pad for finite diffy stencil
    I1, I2 = first(Rc), last(Rc)
    Rc = CartesianIndices(ntuple(ndims(Rc)) do n
        max(1, I1[n]-1) : min(sz[n], I2[n]+1)
    end)

    # TODO: performance test proper cutoff
    szc = size(Rc)
    _crop = maximum(sz .- szc) > 10

    if !_crop
        mc = mask
        d2uc = d2u
        uc = u
    else
        mc = similar(mask, szc)
        mc = _tcopyto!(mc, @view(mask[Rc]))

        if N == 3
            d2uc = @view(d2u[Rc])
            uc = similar(u, szc)
            uc = _tcopyto!(uc, @view(u[Rc]))
        else
            d2uc = @view(d2u[Rc,:])
            uc = similar(u, (szc..., size(u, 4)))
            uc = _tcopyto!(uc, @view(u[Rc,:]))
        end
    end

    # solve
    A = Poisson2(T, mc, dx)
    uc = mgpcg!(uc, A, d2uc; kwargs...)

    if _crop
        if N == 3
            _tcopyto!(@view(u[Rc]), uc)
        else
            _tcopyto!(@view(u[Rc,:]), uc)
        end
    end

    return u
end

function solve_poisson_dct(
    d2u::AbstractArray{<:AbstractFloat, N},
    dx::NTuple{3, Real}
) where {N}
    N ∈ (3, 4) || throw(ArgumentError("arrays must be 3d or 4d, got $(N)d"))
    return solve_poisson_dct!(tcopy(d2u), dx)
end

function solve_poisson_dct!(
    u::AbstractArray{<:AbstractFloat, N},
    d2u::AbstractArray{<:AbstractFloat, N},
    dx::NTuple{3, Real}
) where {N}
    N ∈ (3, 4) || throw(ArgumentError("arrays must be 3d or 4d, got $(N)d"))
    size(u) == size(d2u) || throw(DimensionMismatch())
    return solve_poisson_dct!(_tcopyto!(u, d2u), dx)
end

function solve_poisson_dct!(
    d2u::AbstractArray{<:AbstractFloat, N},
    dx::NTuple{3, Real}
) where {N}
    N ∈ (3, 4) || throw(ArgumentError("arrays must be 3d or 4d, got $(N)d"))

    nx, ny, nz = size(d2u)[1:3]
    idx2 = inv(dx[1].*dx[1])
    idy2 = inv(dx[2].*dx[2])
    idz2 = inv(dx[3].*dx[3])

    u = d2u

    FFTW.set_num_threads(num_cores())
    P = plan_dct!(u, 1:3)
    iP = inv(P)

    u = P*u

    X = [2*(cospi(i)-1)*idx2 for i in range(0, step=1/nx, length=nx)]
    Y = [2*(cospi(j)-1)*idy2 for j in range(0, step=1/ny, length=ny)]
    Z = [2*(cospi(k)-1)*idz2 for k in range(0, step=1/nz, length=nz)]

    @inbounds for t in axes(u, 4)
        d2ût = @view(u[:,:,:,t])
        @batch for k in 1:nz
            for j in 1:ny
                for i in 1:nx
                    d2ût[i,j,k] *= inv(X[i] + Y[j] + Z[k])
                end
            end
        end
        # X = Y = Z = 0. Solution is not unique. Set mean to 0
        d2ût[1,1,1] = 0
    end

    # inverse dct
    return iP*u
end

function solve_poisson_fft(
    d2u::AbstractArray{<:AbstractFloat, N},
    dx::NTuple{3, Real}
) where {N}
    N ∈ (3, 4) || throw(ArgumentError("arrays must be 3d or 4d, got $(N)d"))
    return solve_poisson_fft!(similar(d2u), d2u, dx)
end

function solve_poisson_fft!(
    u::AbstractArray{<:AbstractFloat, N},
    d2u::AbstractArray{<:AbstractFloat, N},
    dx::NTuple{3, Real}
) where {N}
    N ∈ (3, 4) || throw(ArgumentError("arrays must be 3d or 4d, got $(N)d"))
    size(u) == size(d2u) || throw(DimensionMismatch())

    nx, ny, nz = size(d2u)[1:3]
    idx2 = inv(dx[1].*dx[1])
    idy2 = inv(dx[2].*dx[2])
    idz2 = inv(dx[3].*dx[3])

    _rfft = iseven(nx)
    FFTW.set_num_threads(num_cores())

    # FFTW's rfft is extremely slow with some odd lengths in the first dim
    if _rfft
        nx = nx>>1 + 1
        P = plan_rfft(d2u, 1:3)
        iP = inv(P)
        d2û = P*d2u
    else
        d2û = _tcopyto!(similar(d2u, complex(eltype(d2u))), d2u)
        P = plan_fft!(d2û, 1:3)
        iP = inv(P)
        d2û = P*d2û
    end

    X = [2*(cospi(i)-1)*idx2 for i in fftfreq(size(u, 1), 2)]
    Y = [2*(cospi(j)-1)*idy2 for j in fftfreq(size(u, 2), 2)]
    Z = [2*(cospi(k)-1)*idz2 for k in fftfreq(size(u, 3), 2)]

    @inbounds for t in axes(d2û, 4)
        d2ût = @view(d2û[:,:,:,t])
        @batch for k in 1:nz
            for j in 1:ny
                for i in 1:nx
                    d2ût[i,j,k] *= inv(X[i] + Y[j] + Z[k])
                end
            end
        end
        # X = Y = Z = 0. Solution is not unique. Set mean to 0
        d2ût[1,1,1] = 0
    end

    # inverse fft
    if _rfft
        u = mul!(u, iP, d2û)
    else
        d2û = iP*d2û
        u = _tcopyto!(real, u, d2û)
    end

    return u
end


@testset "Points" begin
  @test embeddim(Point(1, 2)) == 2
  @test embeddim(Point(1, 2, 3)) == 3
  @test coordtype(Point(1, 1)) == Int
  @test coordtype(Point(1.,1.)) == Float64
  @test coordtype(Point(1f0, 1f0)) == Float32
  @test coordtype(Point2(1, 1)) == Float64
  @test coordtype(Point3(1, 1, 1)) == Float64
  @test coordtype(Point2f(1, 1)) == Float32
  @test coordtype(Point3f(1, 1, 1)) == Float32

  @test coordtype(Point{2,T}((1, 1))) == T
  @test coordtype(Point{2,T}(1, 1)) == T

  @test coordinates(P2(1, 2)) == [1.0, 2.0]
  @test coordinates(P3(1, 2, 3)) == [1.0, 2.0, 3.0]

  @test P2(1, 2) - P2(1, 1) == T[0, 1]
  @test P3(1, 2, 3) - P3(1, 1, 1) == T[0, 1, 2]
  @test_throws DimensionMismatch P2(1, 2) - P3(1, 2, 3)

  @test P2(1, 2) + V2(0, 0) == P2(1, 2)
  @test P2(2, 3) + V2(2, 1) == P2(4, 4)
  @test P3(1, 2, 3) + V3(0, 0, 0) == P3(1, 2, 3)
  @test P3(2, 3, 4) + V3(2, 1, 0) == P3(4, 4, 4)
  @test_throws DimensionMismatch P2(1, 2) + V3(1, 2, 3)

  @test P2(1, 2) - V2(0, 0) == P2(1, 2)
  @test P2(2, 3) - V2(2, 1) == P2(0, 2)
  @test P3(1, 2, 3) - V3(0, 0, 0) == P3(1, 2, 3)
  @test P3(2, 3, 4) - V3(2, 1, 0) == P3(0, 2, 4)

  @test embeddim(rand(P2)) == 2
  @test embeddim(rand(P3)) == 3
  @test coordtype(rand(P2)) == T
  @test coordtype(rand(P3)) == T

  @test eltype(rand(P2, 3)) == P2
  @test eltype(rand(P3, 3)) == P3

  @test P2(1, 2) ≈ P2(1 + eps(T), T(2))
  @test P3(1, 2, 3) ≈ P3(1 + eps(T), T(2), T(3))

  @test embeddim(Point([1,2])) == 2
  @test coordtype(Point([1,2])) == Int
  @test coordtype(Point([1.,2.])) == Float64

  @test embeddim(Point([1,2,3])) == 3
  @test coordtype(Point([1,2,3])) == Int
  @test coordtype(Point([1.,2.,3.])) == Float64
end


using Reexport

@reexport module TransferOperators

    # Abstract types and composite types for different estimators.
    include("AbstractTransferOperator.jl")

    include("composite_types/TransferOperatorRectangularBinning.jl")
    include("composite_types/TransferOperatorTriangulationExact.jl")
    include("composite_types/TransferOperatorTriangulationApproximate.jl")
    include("composite_types/plot_recipes.jl")

    # Constructors for the estimators.
    include("Estimators/Estimators.jl")

end # module

"""
    TransferOperators

A module containing types and functions to approximate the transfer operator
from numerical data. The estimators in this module are the basis for computing
invariant measures, which is done in the `InvariantMeasures` module.
"""
TransferOperators


function bubblesort!(arr::AbstractVector)
    for _ in 2:length(arr), j in 1:length(arr)-1
        if arr[j] > arr[j+1]
            arr[j], arr[j+1] = arr[j+1], arr[j]
        end
    end
    return arr
end

v = rand(-10:10, 10)
println("# unordered: $v\n -> ordered: ", bubblesort!(v))


module Get_kurtosis

export get_kurtosis

using Distributions, Statistics, Dierckx, SeisIO

"""
    get_kurtosis(data::SeisChannel,kurtsis_tw_sparse::Float64; timewinlength::Float64=60)

    compute kurtosis at each timewindow

# Input:
    - `data::SeisData`    : SeisData from SeisIO
    - `kurtosis_tw_sparse::Float64` : time length of span for kurtosis time window
    - `timewinlength::Float64`  : time window to calculate kurtosis
    kurtosis evaluation following Baillard et al.(2013)
"""
function get_kurtosis(data::SeisChannel, kurtosis_tw_sparse::Float64; timewinlength::Float64=60.0)

    #convert window lengths from seconds to samples
    TimeWin = trunc(Int,timewinlength * data.fs)
    data.misc["kurtosis"] = fast_kurtosis_series(data.x, TimeWin, kurtosis_tw_sparse)

    return data

end

"""
    fast_kurtosis_series(v::RealArray, TimeWin::Int64)

    fast compute kurtosis series at each timewindow

# Input:
    - `v::RealArray`    : SeisData from SeisIO
    - `N::Int64`  : time window length to calculate kurtosis

    kurtosis evaluation following Baillard et al.(2013)
"""
function fast_kurtosis_series(v::Array{Float64, 1}, TN::Int64, kurtosis_tw_sparse::Float64)

    kurt = zeros(length(v))
    n = length(v)
    kurt_grid = 1:n

    KTWSparse = kurtosis_tw_sparse

    if n < TN error("Kurtosis time window is larger than data length. Decrease time window.") end

    cm2 = 0.0  # empirical 2nd centered moment (variance)
    cm4 = 0.0  # empirical 4th centered moment

    # !DEPRECATED!: 1. compute mean value at each time window by numerical sequence
    # 2. use mapreduce to sum values

    # # first term
    # Trace = @views v[1:TN]
    # z2 = zeros(TN)
    # m0 = mean(Trace)
    #
    # cm2 = Statistics.varm(Trace, m0, corrected=false)
    # cm4 = fourthmoment(Trace, m0, corrected=false)
    #
    # # fill first part with kurtosis at TN
    # kurt[1:TN] .= (cm4 / (cm2 * cm2)) - 3.0
    #
    # @simd for k = TN:n-1
    #
    #     diff1 = @inbounds @views (v[k-TN+1] - v[k+1])/TN
    #     m1 = m0 - diff1
    #     Trace = @views v[k-TN+2:k+1]
    #     cm2 = Statistics.varm(Trace, m1, corrected=false)
    #     cm4 = fourthmoment(Trace, m1, corrected=false) #sum(xi - m)^4 / N
    #     kurt[k+1] = (cm4 / (cm2 * cm2)) - 3.0
    #     m0 = m1
    # end

    # first term
    Trace = @views v[1:TN]
    z2 = zeros(TN)
    m0 = mean(Trace)

    cm2 = Statistics.varm(Trace, m0, corrected=false)
    cm4 = fourthmoment(Trace, m0, corrected=false)

    # fill first part with kurtosis at TN
    kurt[1:TN] .= (cm4 / (cm2 * cm2)) - 3.0

    @simd for k = TN:n-1

        diff1 = @inbounds @views (v[k-TN+1] - v[k+1])/TN
        m1 = m0 - diff1
        Trace = @views v[k-TN+2:k+1]
        cm2 = Statistics.varm(Trace, m1, corrected=false)
        cm4 = fourthmoment(Trace, m1, corrected=false) #sum(xi - m)^4 / N
        kurt[k+1] = (cm4 / (cm2 * cm2)) - 3.0
        m0 = m1
    end

    return kurt

end

#---following functions are modified from Statistics.jl---#

centralizedabs4fun(m) = x -> abs2.(abs2.(x - m))
centralize_sumabs4(A::AbstractArray, m) =
    mapreduce(centralizedabs4fun(m), +, A)
centralize_sumabs4(A::AbstractArray, m, ifirst::Int, ilast::Int) =
    Base.mapreduce_impl(centralizedabs4fun(m), +, A, ifirst, ilast)

function centralize_sumabs4!(R::AbstractArray{S}, A::AbstractArray, means::AbstractArray) where S
    # following the implementation of _mapreducedim! at base/reducedim.jl
    lsiz = Base.check_reducedims(R,A)
    isempty(R) || fill!(R, zero(S))
    isempty(A) && return R

    if Base.has_fast_linear_indexing(A) && lsiz > 16 && !has_offset_axes(R, means)
        nslices = div(length(A), lsiz)
        ibase = first(LinearIndices(A))-1
        for i = 1:nslices
            @inbounds R[i] = centralize_sumabs4(A, means[i], ibase+1, ibase+lsiz)
            ibase += lsiz
        end
        return R
    end
    indsAt, indsRt = Base.safe_tail(axes(A)), Base.safe_tail(axes(R)) # handle d=1 manually
    keep, Idefault = Broadcast.shapeindexer(indsRt)
    if Base.reducedim1(R, A)
        i1 = first(Base.axes1(R))
        @inbounds for IA in CartesianIndices(indsAt)
            IR = Broadcast.newindex(IA, keep, Idefault)
            r = R[i1,IR]
            m = means[i1,IR]
            @simd for i in axes(A, 1)
                r += abs2(abs2(A[i,IA] - m))
            end
            R[i1,IR] = r
        end
    else
        @inbounds for IA in CartesianIndices(indsAt)
            IR = Broadcast.newindex(IA, keep, Idefault)
            @simd for i in axes(A, 1)
                R[i,IR] += abs2(abs2(A[i,IA] - means[i,IR]))
            end
        end
    end
    return R
end

function fourthmoment!(R::AbstractArray{S}, A::AbstractArray, m::AbstractArray; corrected::Bool=true) where S
    if isempty(A)
        fill!(R, convert(S, NaN))
    else
        rn = div(length(A), length(R)) - Int(corrected)
        centralize_sumabs4!(R, A, m)
        R .= R .* (1 // rn)
    end
    return R
end

"""
    fourthmoment(v, m; dims, corrected::Bool=true)
Compute the fourthmoment of a collection `v` with known mean(s) `m`,
optionally over the given dimensions. `m` may contain means for each dimension of
`v`. If `corrected` is `true`, then the sum is scaled with `n-1`,
whereas the sum is scaled with `n` if `corrected` is `false` where `n = length(v)`.
!!! note
    If array contains `NaN` or [`missing`](@ref) values, the result is also
    `NaN` or `missing` (`missing` takes precedence if array contains both).
    Use the [`skipmissing`](@ref) function to omit `missing` entries and compute the
    variance of non-missing values.
"""
fourthmoment(A::AbstractArray, m::AbstractArray; corrected::Bool=true, dims=:) = _fourthmoment(A, m, corrected, dims)

_fourthmoment(A::AbstractArray{T}, m, corrected::Bool, region) where {T} =
    fourthmoment!(Base.reducedim_init(t -> abs2(t)/2, +, A, region), A, m; corrected=corrected)

fourthmoment(A::AbstractArray, m; corrected::Bool=true) = _fourthmoment(A, m, corrected, :)

function _fourthmoment(A::AbstractArray{T}, m, corrected::Bool, ::Colon) where T
    n = length(A)
    n == 0 && return oftype((abs2(zero(T)) + abs2(zero(T)))/2, NaN)
    return centralize_sumabs4(A, m) / (n - Int(corrected))
end

end


# This is an extended version of approx_space_time_inference.jl. It combines it with
# Optim + ParameterHandling + Zygote to learn the kernel parameters.
# If you understand how to use Optim + ParameterHandling + Zygote for an AbstractGP,
# e.g. that shown on the README for this package, and how approx_space_time_inference.jl
# works, then you should understand this file.

using AbstractGPs
using KernelFunctions
using TemporalGPs

using TemporalGPs: Separable, approx_posterior_marginals, RegularInTime

# Load standard packages from the Julia ecosystem
using Optim # Standard optimisation algorithms.
using ParameterHandling # Helper functionality for dealing with model parameters.
using Zygote # Algorithmic Differentiation

using ParameterHandling: flatten

# Declare model parameters using `ParameterHandling.jl` types.
flat_initial_params, unflatten = flatten((
    var_kernel = positive(0.6),
    λ_space = positive(0.5),
    λ_time = positive(0.1),
    var_noise = positive(0.1),
));

# Construct a function to unpack flattened parameters and pull out the raw values.
unpack = ParameterHandling.value ∘ unflatten;
params = unpack(flat_initial_params);

function build_gp(params)
    k_space = SEKernel() ∘ ScaleTransform(params.λ_space)
    k_time = Matern52Kernel() ∘ ScaleTransform(params.λ_time)
    k = params.var_kernel * Separable(k_space, k_time)
    return to_sde(GP(k), ArrayStorage(Float64))
end

# Construct inputs. Spatial locations change at each point in time.
# Also works with RectilinearGrids of inputs.
# Times must be increasing, points in space can be anywhere.
N = 50;
T = 1000;
points_in_space = [randn(N) for _ in 1:T];
points_in_time = RegularSpacing(0.0, 0.1, T);
x = RegularInTime(points_in_time, points_in_space);

# Since it's not straightforward to generate samples from this GP at `x`, use a known
# function, under a bit of iid noise.
xs = collect(x);
y = sin.(first.(xs)) .+ cos.(last.(xs)) + sqrt.(params.var_noise) .* randn(length(xs));

# Spatial pseudo-point inputs.
z_r = collect(range(-3.0, 3.0; length=5));

# Specify an objective function for Optim to minimise in terms of x and y.
# We choose the usual negative log marginal likelihood (NLML).
function objective(params)
    f = build_gp(params)
    return -elbo(f(x, params.var_noise), y, z_r)
end

# Optimise using Optim. Takes a little while to compile because Zygote.
training_results = Optim.optimize(
    objective ∘ unpack,
    θ -> only(Zygote.gradient(objective ∘ unpack, θ)),
    flat_initial_params,
    BFGS(
        alphaguess = Optim.LineSearches.InitialStatic(scaled=true),
        linesearch = Optim.LineSearches.BackTracking(),
    ),
    Optim.Options(show_trace = true);
    inplace=false,
);

# Extracting the final values of the parameters.
# Should be close to truth.
final_params = unpack(training_results.minimizer);
# final_params = unpack(flat_initial_params);

# Locations in space at which to make predictions. Assumed to be the same at each point in
# time, but this assumption could easily be relaxed.
N_pr = 150;
x_r_pr = range(-5.0, 5.0; length=N_pr);

# Compute the approximate posterior marginals.
fx_final = build_gp(final_params)(x, final_params.var_noise)
f_post_marginals = approx_posterior_marginals(dtc, fx_final, y, z_r, x_r_pr);
m_post_marginals = mean.(f_post_marginals);
σ_post_marginals = std.(f_post_marginals);

# Visualise the posterior marginals. We don't do this during in CI because it causes
# problems.
if get(ENV, "TESTING", "FALSE") == "FALSE"
    using Plots
    savefig(
        plot(
            heatmap(reshape(m_post_marginals, N_pr, T)),
            heatmap(reshape(σ_post_marginals, N_pr, T));
            layout=(1, 2),
        ),
        "posterior.png",
    );
end


using Test, Random, Flux

@testset "Weights" begin
    @testset "uniform" begin
        @test size(OperatorLearning.cglorot_uniform(128,64,10)) == (128, 64, 10)
        @test eltype(OperatorLearning.cglorot_uniform(128,64,10)) == ComplexF32
    end

    @testset "normal" begin
        @test size(OperatorLearning.cglorot_normal(128,64,10)) == (128, 64, 10)
        @test eltype(OperatorLearning.cglorot_normal(128,64,10)) == ComplexF32
    end
end


#!/usr/bin/env julia

using Luxor

using Test

using Random
Random.seed!(42)

function point_intersect(fname)
    Drawing(1000, 1000, fname)
    origin()
    setopacity(0.4)

    for i in 1:100
        randomhue()
        p1 = Point(rand(-450:450), rand(-450:450))
        p2 = Point(rand(-450:450), rand(-450:450))
        p3 = Point(rand(-450:450), rand(-450:450))
        p4 = Point(rand(-450:450), rand(-450:450))
        flag, intersection_point = intersectionlines(p1, p2, p3, p4, crossingonly=true)
        if flag
            circle(intersection_point, 5, :fill)
            move(p1)
            line(p2)
            strokepath()
            move(p3)
            line(p4)
            strokepath()
        end
    end
    @test finish() == true
    println("...finished test: output in $(fname)")
end

point_intersect("line_intersection.pdf")


# ---
# title: 1572. Matrix Diagonal Sum
# id: problem1572
# author: Tian Jun
# date: 2020-10-31
# difficulty: Easy
# categories: Array
# link: <https://leetcode.com/problems/matrix-diagonal-sum/description/>
# hidden: true
# ---
#
# Given a square matrix `mat`, return the sum of the matrix diagonals.
#
# Only include the sum of all the elements on the primary diagonal and all the
# elements on the secondary diagonal that are not part of the primary diagonal.
#
#
#
# **Example 1:**
#
# ![](https://assets.leetcode.com/uploads/2020/08/14/sample_1911.png)
#
#
#
#     Input: mat = [[ **1** ,2, **3** ],
#                   [4, **5** ,6],
#                   [ **7** ,8, **9** ]]
#     Output: 25
#     Explanation: Diagonals sum: 1 + 5 + 9 + 3 + 7 = 25
#     Notice that element mat[1][1] = 5 is counted only once.
#
#
# **Example 2:**
#
#
#
#     Input: mat = [[ **1** ,1,1, **1** ],
#                   [1, **1** , **1** ,1],
#                   [1, **1** , **1** ,1],
#                   [ **1** ,1,1, **1** ]]
#     Output: 8
#
#
# **Example 3:**
#
#
#
#     Input: mat = [[ **5** ]]
#     Output: 5
#
#
#
#
# **Constraints:**
#
#   * `n == mat.length == mat[i].length`
#   * `1 <= n <= 100`
#   * `1 <= mat[i][j] <= 100`
#
#
## @lc code=start
using LeetCode

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


# Model
include_model("acrobot")

# Horizon
T = 101

# Time step
tf = 5.0
h = tf / (T - 1)

# ul <= u <= uu
ul, uu = control_bounds(model, T, -10.0, 10.0)

# Initial and final states
x1 = [0.0; 0.0; 0.0; 0.0]
xT = [π; 0.0; 0.0; 0.0]
xl, xu = state_bounds(model, T, x1 = x1, xT = xT)

# Objective
obj = quadratic_tracking_objective(
        [t < T ? Diagonal(1.0 * ones(model.n)) : Diagonal(ones(model.n)) for t = 1:T],
        [Diagonal(1.0e-1 * ones(model.m)) for t = 1:T-1],
        [xT for t = 1:T],
        [zeros(model.m) for t = 1:T])

# Problem
prob = trajectory_optimization_problem(model,
           obj,
           T,
           h = h,
           xl = xl,
           xu = xu,
           ul = ul,
           uu = uu)

# Trajectory initialization
x0 = linear_interpolation(x1, xT, T) # linear interpolation on state
u0 = random_controls(model, T, 0.001) # random controls

# Pack trajectories into vector
z0 = pack(x0, u0, prob)

# Solve
@time z, info = solve(prob, copy(z0), tol = 1.0e-6, c_tol = 1.0e-6)

# Visualize
using Plots
x, u = unpack(z, prob)
plot(hcat(x...)', width = 2.0)
plot(hcat(u...)', width = 2.0, linetype = :steppost)

include(joinpath(pwd(), "models/visualize.jl"))
vis = Visualizer()
render(vis)
visualize!(vis, model, x, Δt = h)

# Model-predictive control
include_snopt()

function run_mpc()
    T_mpc = 25
    x_ref = [x...,[x[end] for i = 1:T_mpc]...]
    u_ref = [u...,[u[end] for i = 1:T_mpc]...]
    Q_mpc = Diagonal(100.0 * ones(model.n))
    R_mpc = Diagonal(0.1 * ones(model.m))
    K, P = tvlqr(model, x, u, h, [Q_mpc for t = 1:T], [R_mpc for t = 1:T-1])
    P_mpc = [P..., [P[end] for i = 1:T_mpc]...]

    model_sim = Acrobot{RK3, FixedTime}(4, 1, 4,
        1.0 + 0.05, 0.33, 1.0, 0.5, 1.0 - 0.05, 0.33, 1.0, 0.5, 9.81, 0.1, 0.1)

    shift = 1

    x_hist = [x_ref[shift] + 1.0e-1 * randn(model.n)]
    u_hist = []

    for t = 1:T-1
    	println("T = $t")
    	ul_mpc, uu_mpc = control_bounds(model, T_mpc, -10.0, 10.0)
    	xl_mpc, xu_mpc = state_bounds(model, T_mpc, x1 = x_hist[end])

    	# Objective
    	obj_mpc = quadratic_tracking_objective(
    	        [t < T_mpc ? Q_mpc : P_mpc[T_mpc + shift - 1] for t = 1:T_mpc],
    	        [R_mpc for t = 1:T_mpc-1],
    	        [x_ref[shift + t - 1] for t = 1:T_mpc],
    		[u_ref[shift + t - 1] for t = 1:T_mpc-1])

    	# Problem
    	prob_mpc = trajectory_optimization_problem(model,
    			obj_mpc,
    			T_mpc,
    			h = h,
    			ul = ul_mpc,
    			uu = uu_mpc,
    			xl = xl_mpc,
    			xu = xu_mpc)

    	# Pack trajectories into vector
    	z0_mpc = pack(x_ref[shift .+ (1:T_mpc)], u_ref[shift .+ (1:T_mpc-1)], prob_mpc)

    	# Solve
    	@time z_mpc, info = solve(prob_mpc, copy(z0_mpc),
            nlp = :SNOPT7, mapl = 0)
    	x_mpc, u_mpc = unpack(z_mpc, prob_mpc)

    	push!(x_hist,
    		fd(model_sim, x_hist[end], u_mpc[1], 1.0e-3 * randn(model_sim.d), h, nothing))
    	push!(u_hist, u_mpc[1])
    	shift += 1
    end
    return x_hist, u_hist
end

x_hist_mpc, u_hist_mpc = run_mpc()

plot(hcat(x...)',
	label = "", color = :black, width = 2.0)
plot!(hcat(x_hist_mpc...)',
	label = "", color = :magenta, width = 1.0)

plot(hcat(u...)',
	label = "u (ref)", color = :black, width = 2.0)
plot!(hcat(u_hist_mpc...)',
	label = "u (mpc)", color = :magenta, width = 1.0)

include(joinpath(pwd(), "models/visualize.jl"))
vis = Visualizer()
render(vis)
visualize!(vis, model, x_hist_mpc, Δt = h)


### A Pluto.jl notebook ###
# v0.12.4

using Markdown
using InteractiveUtils

# ╔═╡ b80a0fc0-0675-11eb-1fcd-277d5ec1c289
using Pkg, DrWatson

# ╔═╡ b80a447c-0675-11eb-319c-01dc82d3fb7c
begin
	@quickactivate "SR2TuringPluto"
	using Turing
	using StatisticalRethinking
end

# ╔═╡ b221c3e6-0675-11eb-1ea1-4307128c6de3
md"## Clip-04-07-25t.jl"

# ╔═╡ b80ac79e-0675-11eb-2e70-17f8aa994dc0
md"## Clip-04-97-25t.jl"

# ╔═╡ b81d124e-0675-11eb-16e7-91a1a32f0369
md"### snippets 4.7 - 4.11"

# ╔═╡ b81df8d2-0675-11eb-33ad-819faf448b8a
begin
	df = CSV.read(sr_datadir("Howell1.csv"), DataFrame)
	df = df[df.age .>= 18, :]
end;

# ╔═╡ b82e9b88-0675-11eb-27b5-69985a392610
md"### snippets 4.12, 4.13"

# ╔═╡ b82f4f42-0675-11eb-3f77-d9e53da11c70
begin
	fig1 = plot(p -> pdf(Normal(178, 20), p), 100:250, leg=false, title="prior mu")
	fig2 = plot(p -> pdf(Uniform(0, 50), p), -10:60, leg=false, title="prior sigma")
	plot(fig1, fig2, layout=(1, 2))
end

# ╔═╡ b83efd16-0675-11eb-2e8d-fd33be380806
md"### snippet 4.14"

# ╔═╡ b84a549a-0675-11eb-08c9-b908f9639415
begin
	N = 1000
	d_μ = Normal(178, 20)
	sample_mu = rand(d_μ, N)
	d_σ = Uniform(0, 50)
	sample_sigma = rand(d_σ, N)
	prior_h = rand.(Normal.(sample_mu, sample_sigma))
	density(prior_h)
end

# ╔═╡ b85bf01a-0675-11eb-17e9-67c60998b96b
md"### snippet 4.15"

# ╔═╡ b8664768-0675-11eb-07de-555dbc1de52b
begin
	d_μ_100 = Normal(178, 100)
	sample_μ_100 = rand(d_μ_100, N)
end;

# ╔═╡ b86fd274-0675-11eb-2564-dfb336f2f5e6
begin
	sample_height_100 = rand.(Normal.(sample_μ_100, sample_sigma))
	density(sample_height_100)
end

# ╔═╡ 97fbe892-0676-11eb-0a85-6d793b380f4e
md"### snippets 4.16, 4.17"

# ╔═╡ 9b8b89fe-0676-11eb-1eaa-e986dcba4794
begin
	μ_grid_1 = range(150, 160, length = 50)
	σ_grid_1 = range(7, 9, length = 50)
	post_1 = [(; :μ => μ, :σ => σ) for μ in μ_grid_1, σ in σ_grid_1]  # this is `post`
	ll_1 = [loglikelihood(Normal(μ, σ), df.height) for (μ, σ) in post_1]
	prior_1 = [logpdf(d_μ, μ) + logpdf(d_σ, σ) for (μ, σ) in post_1]
	lposterior_1 = ll_1 .+ prior_1
	lposterior_1 .-= maximum(lposterior_1)
	posterior_1 = exp.(lposterior_1)
	wireframe(μ_grid_1, σ_grid_1, posterior_1)
end

# ╔═╡ b897b90a-0676-11eb-1fe9-cbd663edb4f0
md"### snippets 4.19 - 4.22"

# ╔═╡ c3a4dba2-0676-11eb-07fc-bfb72963914b
begin
	samples = sample(post_1, pweights(posterior_1), N) |> DataFrame
	scatter(samples.μ, samples.σ, alpha = 0.05, ms = 5, xlabel = "μ", ylabel = "σ")
end

# ╔═╡ fdacb630-0676-11eb-0217-f1e07c127aa2
begin
	fig3 = density(samples.μ, leg=false, title="samples.mu")
	fig4 = density(samples.σ, leg=false, title="samples.sigma")
	plot(fig3, fig4, layout=(1,2))
end

# ╔═╡ 643e622c-0677-11eb-187b-033fe474252d
quantile(samples.μ, (0.1, 0.9))

# ╔═╡ 6eed8892-0677-11eb-3c04-e9cb30ea2b66
quantile(samples.σ, (0.1, 0.9))

# ╔═╡ 776436b0-0677-11eb-19bf-bd953d980210
md"### snippets 4.23 - 4.25"

# ╔═╡ 9ba86ca0-0677-11eb-3304-118f06fda2f0
begin
	d3_height = df.height[1:20]
	μ_grid = range(140, 160, length = 50)
	σ_grid = range(4, 20, length = 50)
	post = [(μ = μ, σ = σ) for μ in μ_grid, σ in σ_grid]
	ll = [loglikelihood(Normal(μ, σ), d3_height) for (μ, σ) in post]
	prior = [logpdf(d_μ, μ) + logpdf(d_σ, σ) for (μ, σ) in post]
	lposterior = ll .+ prior
	lposterior .-= maximum(lposterior)
	posterior = exp.(lposterior)
	wireframe(μ_grid, σ_grid, posterior)
end

# ╔═╡ 8613457c-0675-11eb-10bd-217d514940ef
begin
	samples2 = sample(post, pweights(posterior), N) |> DataFrame
	scatter(samples2.μ, samples2.σ, alpha = 0.05, ms = 5)
end

# ╔═╡ d8ab7426-067a-11eb-0bac-1bb4747061b1
begin
	fig5 = density(samples2.μ, leg=false, title="samples2.mu")
	fig6 = density(samples2.σ, leg=false, title="samples2.sigma")
	plot(fig1, fig2, fig3, fig4, fig5, fig6, layout=(3,2))
end

# ╔═╡ d8d3d31c-067a-11eb-2a92-ef47acab4194
md"## End of clip-04-07-25t.jl"

# ╔═╡ Cell order:
# ╠═b221c3e6-0675-11eb-1ea1-4307128c6de3
# ╠═b80a0fc0-0675-11eb-1fcd-277d5ec1c289
# ╠═b80a447c-0675-11eb-319c-01dc82d3fb7c
# ╟─b80ac79e-0675-11eb-2e70-17f8aa994dc0
# ╟─b81d124e-0675-11eb-16e7-91a1a32f0369
# ╠═b81df8d2-0675-11eb-33ad-819faf448b8a
# ╟─b82e9b88-0675-11eb-27b5-69985a392610
# ╠═b82f4f42-0675-11eb-3f77-d9e53da11c70
# ╟─b83efd16-0675-11eb-2e8d-fd33be380806
# ╠═b84a549a-0675-11eb-08c9-b908f9639415
# ╟─b85bf01a-0675-11eb-17e9-67c60998b96b
# ╠═b8664768-0675-11eb-07de-555dbc1de52b
# ╠═b86fd274-0675-11eb-2564-dfb336f2f5e6
# ╟─97fbe892-0676-11eb-0a85-6d793b380f4e
# ╠═9b8b89fe-0676-11eb-1eaa-e986dcba4794
# ╟─b897b90a-0676-11eb-1fe9-cbd663edb4f0
# ╠═c3a4dba2-0676-11eb-07fc-bfb72963914b
# ╠═fdacb630-0676-11eb-0217-f1e07c127aa2
# ╠═643e622c-0677-11eb-187b-033fe474252d
# ╠═6eed8892-0677-11eb-3c04-e9cb30ea2b66
# ╟─776436b0-0677-11eb-19bf-bd953d980210
# ╠═9ba86ca0-0677-11eb-3304-118f06fda2f0
# ╠═8613457c-0675-11eb-10bd-217d514940ef
# ╠═d8ab7426-067a-11eb-0bac-1bb4747061b1
# ╟─d8d3d31c-067a-11eb-2a92-ef47acab4194


function projection(d::Integer, q::Integer, elType=typeof(1.0))
    # Proj(deriv) = kron(diagm(0 => ones(elType, d)), [i==(deriv+1) ? 1 : 0 for i in 1:q+1]')

    # Slightly faster version of the above:
    D = d * (q + 1)
    Proj(deriv) = begin
        P = zeros(elType, d, D)
        @simd ivdep for i in deriv*d+1:D+1:d*D
            @inbounds P[i] = 1
        end
        return P
    end
    return Proj
end


using LinearAlgebra
using Test

using UnitBlockAngular

@testset "UnitBlockAngular" begin

Aref = [
    1.0 0.0 1.1 1.2 1.3 1.4;
    0.0 1.0 2.1 2.2 2.3 2.4;
    1.0 0.0 3.1 3.2 3.3 3.4;
    0.0 0.0 1.0 0.0 1.0 0.0;
    0.0 0.0 0.0 1.0 0.0 1.0
]

m0, n0, n, R = 3, 2, 4, 2
B0 = Aref[1:3, 1:2]
B  = Aref[1:3, 3:end]
blocks = [1, 2, 1, 2]

A = UnitBlockAngularMatrix{Float64}(3, 2, 4, 2);
A.B0 .= B0
A.B  .= B
A.blockidx .= blocks

@test size(A) == (m0+R, n0+n)
@test A == Aref

A = UnitBlockAngularMatrix(B0, B, R, blocks)
@test size(A) == (m0+R, n0+n)
@test A == Aref

x, xref = ones(n0 + n), ones(n0 + n)
y, yref = zeros(m0 + R), zeros(m0 + R)

@testset "Matrix-vector" begin
    mul!(yref, Aref, xref)
    mul!(y, A, x)
    @test y ≈ yref

    mul!(yref, Aref, xref, 100.0, 10.0)
    mul!(y, A, x, 100.0, 10.0)
    @test y ≈ yref

    y .= 1
    yref .= 1
    mul!(xref, Aref', yref)
    mul!(x, A', y)
    @test x ≈ xref

    y .= 1
    yref .= 1
    mul!(xref, Aref', yref, 100.0, 10.0)
    mul!(x, A', y, 100.0, 10.0)
    @test y ≈ yref
end

end  # testset

# Factorization
include("kkt_solver.jl")


# Day 20: Trench Map
# Image is permuted for faster iteration
function load(::Day{20}, path)
    # function load(path)
    lines = readlines(path)
    alg = [c == '#' for c in lines[1]]
    img = BitMatrix(reduce(hcat, [[c == '#' for c in line] for line in lines[3:end]]))
    return alg, img
end

# Part 1 – count pixels in output image
# Offsets defining neighborhood. We need to "read across each column":
const CONV_OFFSETS = [
    CartesianIndex(-1, -1),
    CartesianIndex(0, -1),
    CartesianIndex(1, -1),
    CartesianIndex(-1, 0),
    CartesianIndex(0, 0),
    CartesianIndex(1, 0),
    CartesianIndex(-1, 1),
    CartesianIndex(0, 1),
    CartesianIndex(1, 1),
]

function pad_image(img; px=1, val=false)
    h, w = size(img)
    padrows = zeros(Bool, px, w)
    padcols = zeros(Bool, h + 2 * px, px)
    return hcat(padcols, vcat(padrows, img, padrows), padcols)
end

@inline getval(img, I::CartesianIndex, outer) = checkbounds(Bool, img, I) ? img[I] : outer

const POWS = 2 .^ collect(8:-1:0)
@inline neighbors2num(n) = sum(POWS .* n) + 1

function apply_algorithm!(out, img, outer, alg)
    @inbounds @simd for I in CartesianIndices(img) # not sure @simd does anything here...
        # Index the value of the algorithm `alg` based on the binary number
        # that results from the neighboring pixels.
        out[I] = alg[neighbors2num([
            getval(img, I + offset, outer) for offset in CONV_OFFSETS
        ])]
    end
    return out
end

function enhance_n_times(alg, img, n)
    img = pad_image(img; px=n)
    buff = similar(img)
    outer = false

    for _ in 1:n
        img .= apply_algorithm!(buff, img, outer, alg)
        # The "outer values" are flashing if `alg[1] = 1` and `alg[end] = 0`
        alg[1] == 1 && (outer = !outer)
    end
    return img
end

solve1(::Day{20}, data) = sum(enhance_n_times(data..., 2))

# Part 2 – Part 1 with more steps
solve2(::Day{20}, data) = sum(enhance_n_times(data..., 50))

testresult1(::Day{20}) = 35
testresult2(::Day{20}) = 3351


# Deltamax Estimator
# Estimate the number of servers using the Deltamax algorithm
# Input: arrival order sorted by departure order
# Output: estimated number of servers
function c_deltamax(out_order::Vector{Int64})
  N = size(out_order, 1)
  currmax = out_order[1]
  deltamax = 0
  for i = 2:N
    nextmax = max(out_order[i], currmax)
    deltamax = max(deltamax, nextmax - currmax)
    currmax = nextmax
  end
  deltamax
end

# Order-based Estimator (Sets)
# Estimate the number of servers using the Order-based algorithm
# Input: arrival order sorted by departure order
# Output: estimated number of servers
function c_order_slow(out_order::Vector{Int64})
    N = size(out_order, 1)
    currmax = 0
    nservers = 0
    for i = 1:N
        currmax = max(out_order[i], currmax)
        custall = Set(1:currmax)
        custdeparted = Set(out_order[1:i])
        custservice = setdiff(custall, custdeparted)
        nservers = max(length(custservice) + 1, nservers)
    end
    nservers
end

# Order-based Estimator (Scalar)
# Estimate the number of servers using the Order-based algorithm
# Input: arrival order sorted by departure order
# Output: estimated number of servers
function c_order(out_order::Vector{Int64})
    N = size(out_order, 1)
    currmax = 0
    noccupied = 0
    nservers = 0
    for i = 1:N
        nextmax = max(out_order[i], currmax)
        if out_order[i] < currmax
            noccupied = noccupied - 1
        else
            noccupied = noccupied + nextmax - currmax - 1
        end
        currmax = nextmax
        nservers = max(nservers, noccupied + 1)
    end
    nservers
end

# WORK IN PROGRESS: Summary Estimator - fast version
# WARNING: PRODUCES INCORRECT OUTPUT
# function c_var_unf_fast(A::Vector{Float64}, D::Vector{Float64}, cmax::Int64)
#   # Initilaize data structures and functions
#   N = size(D, 1)
#   B̂ = zeros(N)
#   temp = zeros(N)
#   VS = zeros(cmax, 2) # var, unf
#   VS[1,:] = [-1,-1]
#   max_undef = 1
#   Dsort = sort(D)
#   ## Grid Search
#   for c = 2:cmax
#     B̂[1:c] = A[1:c]
#     for i = (c + 1):N
#       B̂[i] = max(A[i], Dsort[i - c])
#     end
#     Ŝ = D - B̂
#     if any(Ŝ .< 0)
#       VS[c, 1:2] = -1
#       max_undef = c
#     else
#       VS[c, 1] = var(Ŝ) * (N-1) # var
#     end
#   end
#   (r1, r2) = (99, 99) # initialize result to 99 for debugging
#   # var method
#   try
#     r1 = findin(VS[:, 1], minimum(filter(x -> x >= 0, VS[:, 1])))[1]
#   catch
#     r1 = 1
#   end
#   # uninformed method
#   r2 = max_undef + 1
#   (r1, r2, VS)
# end

# Summary Estimator
# Estimate the number of servers using the Variance and Uninformed methods
# Input: arrival times, departure times, and the max number of servers
# Output: estimated number of servers for each method
function c_var_unf(A::Vector{Float64}, D::Vector{Float64}, cmax::Int64)
  # Initilaize data structures and functions
  N = size(D, 1)
  B̂ = zeros(N)
  temp = zeros(N)
  VS = zeros(cmax, 2) # var, unf
  max_undef = 0
  ## Grid Search
  for c = 1:cmax
    B̂[1:c] = A[1:c]
    for i = (c + 1):N
      Dₖ = sort(D[1:(i - 1)])
      B̂[i] = max(A[i], Dₖ[i - c])
    end
    Ŝ = D - B̂
    if any(Ŝ .< 0)
      VS[c, 1:2] = -1
      max_undef = c
    else
      VS[c, 1] = var(Ŝ) * (N-1) # var
    end
  end
  (r1, r2) = (99, 99) # initialize result to 99 for debugging
  # var method
  try
    r1 = findin(VS[:, 1], minimum(filter(x -> x >= 0, VS[:, 1])))[1]
  catch
    r1 = 1
  end
  # uninformed method
  r2 = max_undef + 1
  (r1, r2, VS)
end


module ch12
using Printf
# ------
# Tuples
# https://benlauwens.github.io/ThinkJulia.jl/latest/book.html#chap12
# ------

# --------------------
# Tuples Are Immutable
# https://benlauwens.github.io/ThinkJulia.jl/latest/book.html#_tuples_are_immutable
# --------------------
x = ('a', 'b')
typeof(x)

# single element tuple
x = ('a')
typeof(x) # Char
x = ('a', )
typeof(x) # Tuple{Char}

# tuple()
x = tuple(1, 'a', pi)
x[2]
x[2:3]

# compare
y = tuple(1, 'a', pi)

x == y
x > y

# ----------------
# Tuple Assignment
# ----------------
a, b = 1, 2, 3

a = 5
b = 7
a, b = b, a

# split
addr = "julius.caesar@rome"
uname, domain = split(addr, '@')

# -----------------------
# Tuples as Return Values
# -----------------------
x = divrem(7, 3)
q,r = divrem(7,3)
# -------------------------------
# Variable-length Argument Tuples
# -------------------------------

# gather (slurp)
function printall(args...)
    println(args)
end
printall(1, 2, '3')

# scatter (splat)
x = (7,3)
# divrem(x) error
divrem(x...)

# Exercise 12-1
function sumall(args...)
    sum = 0
    for value in args
        sum += value
    end

    return sum
end

sumall(1,2,3)

# -----------------
# Arrays and Tuples
# -----------------
str = "abc"
x = [1, 2, 3]
# zip
z = zip(str, x)
for p in z
    println(p)
end
#println(collect(z)) # [('a', 1), ('b', 2), ('c', 3)]

for (char, num) in z
    @printf("%g %s\n", num, char)
end

# enumerate
for (ind, char) in enumerate("abc")
    #println(typeof(ind))
    #println(typeof(char))
    @printf("%i %s\n", ind, char)
end

# -----------------------
# Dictionaries and Tuples
# -----------------------
dict = Dict('a'=>1, 'b'=>2, 'c'=>3)
for (key,value) in dict
    @printf("%s %i\n", key, value)
end

z = zip(str, x)
dict = Dict(z)
for (key,value) in dict
    @printf("%s %i\n", key, value)
end

dict = Dict()
dict['a', 1] = 3
dict['b', 1] = 4
for ((well, stage), num) in dict
    @printf("%s,%i: %i\n", well, stage, num)
end

# ----------------------
# Sequences of Sequences
# ----------------------
println(sort([13, 7, 5]))
println(reverse([13, 7, 5]))
end


using PlotlyJS
Random.seed!(42)

function heatmap1()
    plot(heatmap(z=[1 20 30; 20 1 60; 30 60 1]))
end

function heatmap2()
    trace = heatmap(
        x=["Monday", "Tuesday", "Wednesday", "Thursday", "Friday"],
        y=["Morning", "Afternoon", "Evening"],
        z=rand(1:30, 5, 3)
    )
    plot(trace)
end


const defcolors = ["#1F77B4", "#FF7F0E", "#2CA02C", "#D62728", "#9467BD",
                   "#8C564B", "#E377C2", "#7F7F7F", "#BCBD22", "#17BECF"]

sv = IScatterSpectrum.ScatterVolume(230e6, 1e-6, 50000e-9, 0.0)
p  = IScatterSpectrum.Plasma(1.5e11, 2000., 1000., sv)
f  = 5.0:5:5000
freq = [0; f]
tr   = GenericTrace[]
push!(tr, scatter(x=freq,
                  y=[IScatterSpectrum.pwrspec(p, sv);
                     [IScatterSpectrum.pwrspec(freq, p, sv) for freq in f]],
                  line_color=defcolors[1],
                  name="Te=2000 K"))
p  = IScatterSpectrum.Plasma(1.5e11, 1000., 1000., sv)
push!(tr, scatter(x=freq,
                  y=[IScatterSpectrum.pwrspec(p, sv);
                     [IScatterSpectrum.pwrspec(freq, p, sv) for freq in f]],
                  line_color=defcolors[2],
                  name="Te=1000 K"))
p  = IScatterSpectrum.Plasma(1.5e11, 1500., 1000., sv)
push!(tr, scatter(x=freq,
                  y=[IScatterSpectrum.pwrspec(p, sv);
                     [IScatterSpectrum.pwrspec(freq, p, sv) for freq in f]],
                  line_color=defcolors[3],
                  name="Te=1500 K"))
p  = IScatterSpectrum.Plasma(1.5e11, 3000., 1000., sv)
push!(tr, scatter(x=freq,
                  y=[IScatterSpectrum.pwrspec(p, sv);
                     [IScatterSpectrum.pwrspec(freq, p, sv) for freq in f]],
                  line_color=defcolors[4],
                  name="Te=3000 K"))
layout = Layout(xaxis_title="Frequency in Hz",
                yaxis_title="Power",
                font_size=18,
                legend_x=0.85,
                title="Collisionless Incoherent Scatter Spectrum")
plt = plot(tr, layout)
PlotlyJS.savefig(plt, "/home/scb/code/julia/pkg/v0.6/IScatterSpectrum/test/CollisionlessPlasma.pdf")
PlotlyJS.savefig(plt, "/home/scb/code/julia/pkg/v0.6/IScatterSpectrum/test/CollisionlessPlasma.png")
tr2 = [tr[k] for k in [1, 2, 4]]
sv = IScatterSpectrum.ScatterVolume(230e6, 1e-6, 50000e-9, 70*π/180)
p  = IScatterSpectrum.Plasma(1.5e11, 2000., 1000., sv)
push!(tr2, scatter(x=freq,
                   y=[IScatterSpectrum.pwrspec(p, sv);
                      [IScatterSpectrum.pwrspec(freq, p, sv) for freq in f]],
                   line_color=defcolors[1],
                   line_dash="dash",
                   name="α=70 deg"))
p  = IScatterSpectrum.Plasma(1.5e11, 1000., 1000., sv)
push!(tr2, scatter(x=freq,
                   y=[IScatterSpectrum.pwrspec(p, sv);
                      [IScatterSpectrum.pwrspec(freq, p, sv) for freq in f]],
                   line_color=defcolors[2],
                   line_dash="dash",
                   showlegend=false))
p  = IScatterSpectrum.Plasma(1.5e11, 3000., 1000., sv)
push!(tr2, scatter(x=freq,
                   y=[IScatterSpectrum.pwrspec(p, sv);
                      [IScatterSpectrum.pwrspec(freq, p, sv) for freq in f]],
                   line_color=defcolors[4],
                   line_dash="dash",
                   showlegend=false))
sv = IScatterSpectrum.ScatterVolume(230e6, 1e-6, 50000e-9, 30*π/180)
p  = IScatterSpectrum.Plasma(1.5e11, 2000., 1000., sv)
push!(tr2, scatter(x=freq,
                   y=[IScatterSpectrum.pwrspec(p, sv);
                      [IScatterSpectrum.pwrspec(freq, p, sv) for freq in f]],
                   line_color=defcolors[1],
                   line_dash="dot",
                   name="α=30 deg"))
p  = IScatterSpectrum.Plasma(1.5e11, 1000., 1000., sv)
push!(tr2, scatter(x=freq,
                   y=[IScatterSpectrum.pwrspec(p, sv);
                      [IScatterSpectrum.pwrspec(freq, p, sv) for freq in f]],
                   line_color=defcolors[2],
                   line_dash="dot",
                   showlegend=false))
p  = IScatterSpectrum.Plasma(1.5e11, 3000., 1000., sv)
push!(tr2, scatter(x=freq,
                   y=[IScatterSpectrum.pwrspec(p, sv);
                      [IScatterSpectrum.pwrspec(freq, p, sv) for freq in f]],
                   line_color=defcolors[4],
                   line_dash="dot",
                   showlegend=false))
plt = plot(tr2, layout)
PlotlyJS.savefig(plt, "/home/scb/code/julia/pkg/v0.6/IScatterSpectrum/test/CollisionlessPlasmaMagField.pdf")
PlotlyJS.savefig(plt, "/home/scb/code/julia/pkg/v0.6/IScatterSpectrum/test/CollisionlessPlasmaMagField.png")


using Match

import Base: convert

"""
Specifies a particular electronic level within a species.

That is, all the quantum numbers to specify a state, except for the projection
of the total angular momentum on the quantization axis.
"""
abstract type LevelSpec end

"""
Specifies a particular electronic state using the relevant low-field
quantum numbers.

The total spin quantum number is implicitly assumed to be :math:`1 / 2` (cf.
`NoHyperfineOneElectronSpecies`).
"""
immutable NoHyperfineNumberSpec <: LevelSpec
    "Orbital angular momentum quantum number."
    l::Integer

    "Total angular momentum quantum number."
    j::Rational
end

"""
Specifies a particular electronic state using spectroscopic notation.
"""
immutable SpectroscopicSpec <: LevelSpec
    val::String
end

convert{T<:LevelSpec}(::Type{T}, s::String) = convert(T, SpectroscopicSpec(s))

function convert(::Type{NoHyperfineNumberSpec}, s::SpectroscopicSpec)
    val = s.val

    l_sym = val[1]
    l = findfirst(['S', 'P', 'D', 'F', 'G'], l_sym) - 1
    if l == -1
        throw(ArgumentError("Unknown symbol '$(s[1])' for orbital angular momentum quantum number: $(s.val)"))
    end

    val = lstrip(val[2 : end], '_')

    if val[1] == '{'
        if val[end] != '}'
            throw(ArgumentError("Invalid total angular momentum quantum number specification '$val': $(s.val)"))
        end
        val = val[2 : end - 1]
    end

    # Parse numerator.
    fraction_idx = findfirst(val, '/')
    if fraction_idx < 2
        throw(ArgumentError("Expected fractional total angular momentum quantum number, not '$val': $(s.val)"))
    end

    j_num_str = val[1 : fraction_idx - 1]
    j_num = try
        parse(Int64, j_num_str)
    catch
        throw(ArgumentError("Invalid numerator '$j_num_str' for total angular momentum quantum number: $(s.val)"))
    end

    # Skip fraction slash.
    val = val[fraction_idx + 1 : end]

    if length(val) == 0
        throw(ArgumentError("Expected denominator for total angular momentum quantum number: $(s.val)"))
    end

    # Skip optional second slash (Julia's fraction notation).
    val = lstrip(val, '/')

    # Parse denominator.
    j_den = try
        parse(Int64, val)
    catch
        throw(ArgumentError("Invalid denominator '$val' for total angular momentum quantum number: $(s.val)"))
    end

    NoHyperfineNumberSpec(l, j_num // j_den)
end

immutable StateSpec
    level::LevelSpec

    """
    Total momentum projection on the quantisation (z) axis (usually denoted
    :math:`m_J` or :math:`m_F` depending on the species).
    """
    m::Rational
end

export LevelSpec, NoHyperfineNumberSpec, convert, StateSpec, SpectroscopicSpec


using Test, LinearMaps, LinearAlgebra, BlockArrays

@testset "Non-traditional axes" begin

    A = rand(ComplexF64,2,4)
    B = PseudoBlockMatrix{ComplexF64}(undef, [2,3], [3,4])

    ax1 = axes(B)[1]
    ax2 = axes(B)[2]
    fill!(B,0)
    B[Block(1),Block(2)] .= A

    N = @inferred LinearMap(B)
    @test axes(N) == (ax1,ax2)
    @test axes(N, 1) == ax1
    @test axes(N, 2) == ax2
    @test_throws ErrorException axes(N, 3)

    @test eltype(N) == eltype(B)

    u = similar(Array{ComplexF64}, ax2)
    v = PseudoBlockVector{ComplexF64}(undef, [3,5])
    w = similar(Array{ComplexF64}, ax1)

    for i in eachindex(u) u[i] = rand(ComplexF64) end
    for i in eachindex(v) v[i] = rand(ComplexF64) end
    for i in eachindex(w) w[i] = rand(ComplexF64) end

    @test B*u == N*u
    @test_throws DimensionMismatch N*v

    # Lu = L*u
    Nu = N*u

    @test axes(Nu)[1] == axes(N)[1] == axes(B)[1]
    @test blocklengths(axes(Nu)[1]) == blocklengths(axes(N)[1]) == blocklengths(axes(B)[1]) == [2,3]

    for trans in (adjoint, transpose)
        Nt = trans(LinearMap(N))
        Bt = trans(B)
        Ntw = Nt*w
        Btw = Bt*w
        @test Ntw ≈ Btw
        @test axes(Ntw)[1] == axes(Nt)[1] == axes(Bt)[1]
        @test blocklengths(axes(Ntw)[1]) == blocklengths(axes(Nt)[1]) == blocklengths(axes(Bt)[1]) == [3,4]
    end

    C = B + 2N
    @test axes(C) === axes(B) === axes(N)
    @test C*u ≈ 3*Nu

    Cu = C*u
    @test axes(C)[1] == ax1
    @test blocklengths(axes(C)[1]) == blocklengths(ax1)

    A = rand(ComplexF64,2,2)
    B = PseudoBlockMatrix{ComplexF64}(undef, [2,2], [2,2])
    ax1 = axes(B)[1]
    ax2 = axes(B)[2]
    fill!(B,0)
    B[Block(1),Block(2)] .= A
    L = LinearMap(B)
    L2 = L*L
    B2 = B*B
    @test axes(L2) == axes(B2)
    B2 = B*Matrix(B)
    L2 = L*LinearMap(Matrix(B))
    @test axes(L2) == axes(B2)
    u = similar(Array{ComplexF64}, ax2)
    B2u = B2*u; L2u = L2*u
    @test axes(B2u)[1] == axes(L2u)[1] == axes(B2)[1] == axes(L2)[1]
    @test blocklengths(axes(B2u)[1]) == blocklengths(axes(L2u)[1]) == [2,2]
end


import DelayDiffEq: MethodOfSteps, Tsit5

@testset "test_delay_system" begin

# For simplicity, equality of DelayLtiSystems are tested over a finite set of frequencies
ω = 0.0:8

@test typeof(promote(delay(0.2), ss(1))[1]) == DelayLtiSystem{Float64,Float64}

@test typeof(promote(delay(0.2), ss(1.0 + im))[1]) == DelayLtiSystem{Complex{Float64}, Float64}

@test sprint(show, ss(1,1,1,1)*delay(1.0)) == "DelayLtiSystem{Float64, Float64}\n\nP: StateSpace{Continuous, Float64}\nA = \n 1.0\nB = \n 0.0  1.0\nC = \n 1.0\n 0.0\nD = \n 0.0  1.0\n 1.0  0.0\n\nContinuous-time state-space model\n\nDelays: [1.0]"

# Extremely basic tests
@test freqresp(delay(1), ω) ≈ reshape(exp.(-im*ω), length(ω), 1, 1) rtol=1e-15
@test freqresp(delay(2.5), ω)[:] ≈ exp.(-2.5im*ω) rtol=1e-15
@test freqresp(3.5*delay(2.5), ω)[:] ≈ 3.5*exp.(-2.5im*ω) rtol=1e-15
@test freqresp(delay(2.5)*1.5, ω)[:] ≈ exp.(-2.5im*ω)*1.5 rtol=1e-15

# Addition of constant
@test evalfr(1 + delay(1.0), 0)[] ≈ 2
@test evalfr(1 - delay(1.0), 0)[] ≈ 0
@test evalfr([2 -delay(1.0)], 0) ≈ [2 -1]

# Stritcly proper system
P1 = DelayLtiSystem(ss(-1.0, 1, 1, 0))
P1_fr = 1 ./ (im*ω .+ 1)
@test freqresp(P1, ω)[:] == P1_fr

# Not stritcly proper system
P2 = DelayLtiSystem(ss(-2.0, -1, 1, 1)) # (s+1)/(s+2)
P2_fr = (im*ω .+ 1) ./ (im*ω .+ 2)
@test freqresp(P2, ω)[:] ≈ P2_fr rtol=1e-15

## Addition
@test freqresp(1 + delay(1), ω)[:] ≈ 1 .+ exp.(-im*ω)
@test freqresp(P1 + delay(1), ω)[:] ≈ P1_fr .+ exp.(-im*ω)

# Substraction
@test freqresp(P1 - delay(1), ω)[:] ≈ P1_fr .- exp.(-im*ω)

## Multiplication by scalar
@test freqresp(2*delay(1), ω)[:] ≈ 2*exp.(-im*ω)
@test freqresp(2*(P1*delay(1)), ω)[:] ≈ 2*P1_fr.*exp.(-im*ω)

## Multiplication
@test freqresp(P1 * delay(1), ω)[:] ≈ P1_fr .* exp.(-im*ω) rtol=1e-15
@test freqresp(delay(1) * P1, ω)[:] ≈ P1_fr .* exp.(-im*ω) rtol=1e-15

@test freqresp(P2 * delay(1), ω)[:] ≈ P2_fr .* exp.(-im*ω) rtol=1e-15
@test freqresp(delay(1) * P2, ω)[:] ≈ P2_fr .* exp.(-im*ω) rtol=1e-15

# Equality
@test P1 == deepcopy(P1)
@test P1 != deepcopy(P2)

# evalfr
s_vec = [0, 1im, 1, 1 + 1im]
@test [evalfr(delay(2), s)[1] for s in s_vec] ≈ [exp(-2*s) for s in s_vec] rtol=1e-16

## Feedback
# The first tests don't include delays, but the linear system is of DelayLtiForm type
# (very simple system so easy to troubleshoot)
@test freqresp(feedback(1.0, P1), ω)[:] ≈ 1 ./ (1 .+  P1_fr) rtol=1e-15
@test freqresp(feedback(P1, 1.0), ω)[:] ≈ P1_fr ./ (1 .+  P1_fr) rtol=1e-15
@test freqresp(feedback(P1, P1), ω)[:] ≈ P1_fr ./ (1 .+ P1_fr .* P1_fr) rtol=1e-15

@test freqresp(feedback(1.0, DelayLtiSystem(ss(0.5))), [0])[:] == [2/3]
@test freqresp(feedback(1.0, P2), ω)[:] ≈ 1 ./ (1 .+ P2_fr)

@test freqresp(feedback(0.5, delay(2.0)), ω) ≈ 0.5 ./ (1 .+ 0.5*exp.(-2im*ω))
@test freqresp(feedback(delay(2.0), 0.5), ω) ≈ exp.(-2im*ω) ./ (1 .+ 0.5*exp.(-2im*ω))

@test freqresp(feedback(P1, delay(1)), ω)[:] ≈ P1_fr ./ (1 .+ exp.(-im*ω) .* P1_fr) rtol=1e-15
@test freqresp(feedback(delay(1), P1), ω)[:] ≈ exp.(-im*ω) ./ (1 .+ exp.(-im*ω) .* P1_fr) rtol=1e-15 #FIXME: Answer is Inf, but should give error.. rtol=1e-15
@test freqresp(feedback(P1*delay(1), 1.0), ω)[:] ≈ P1_fr .* exp.(-im*ω) ./ (1 .+ exp.(-im*ω) .* P1_fr) rtol=1e-15
@test freqresp(feedback(1.0, P1*delay(1)), ω)[:] ≈ 1 ./ (1 .+ exp.(-im*ω) .* P1_fr) rtol=1e-15

@test freqresp(feedback(1.0, P2*0.5*(ss(1.0) + delay(2))), ω)[:] ≈ 1 ./(1 .+ P2_fr .* 0.5.*(1 .+ exp.(-2*im*ω)))

@test freqresp(1.0 + delay(2), ω)[:] ≈ 1 .+ exp.(-2im*ω) rtol=1e-15

G = ss(0.5) + 0.5*delay(2)# + 0.5*delay(3)
G_fr = 0.5 .+ 0.5*exp.(-2*im*ω)# + 0.5*exp.(-3*im*ω)
@test freqresp(G, ω)[:] ≈ G_fr rtol=1e-15

@test freqresp(feedback(1.0, P1*G), ω)[:] ≈ 1 ./(1 .+ P1_fr .* G_fr) rtol=1e-15
@test freqresp(feedback(P1, G), ω)[:] ≈ P1_fr ./(1 .+ P1_fr .* G_fr) rtol=1e-15

@test freqresp(feedback(P2, G), ω)[:] ≈ P2_fr ./(1 .+ P2_fr .* G_fr) rtol=1e-15

s = tf("s")

# Test alternative exp constructor for delays
d = exp(-2*s)
@test freqresp(d, [0, 1, 2]) ≈ [1, exp(-2im), exp(-4im)]

@test_throws ErrorException exp(-s^2 - 2*s)
@test_throws ErrorException exp(-2*s+1) # in principle ok, but not allowed anyway
@test_throws ErrorException exp([-2*s; -s])
@test_throws ErrorException exp(2*s) # Non-causal

Ω = [0.5, 1.0, 1.5, 2.0]
# Test for internal function delayd_ss
@test freqresp(ControlSystems.delayd_ss(1.0, 0.2), Ω)[:] ≈ exp.(-im*Ω) atol=1e-14
@test freqresp(ControlSystems.delayd_ss(3.2, 0.4), Ω)[:] ≈ exp.(-3.2*im*Ω) atol=1e-14
@test_throws ErrorException freqresp(ControlSystems.delayd_ss(3.2, 0.5), Ω)

# Simple tests for c2d of DelayLtiSystems
@test freqresp(c2d(feedback(ss(0,1,1,0), delay(1.5)), 0.5), Ω) ≈ [0.5/((z - 1) + 0.5*z^-3) for z in exp.(im*Ω*0.5)]
@test freqresp(c2d(feedback(delay(1.5), delay(1.0)), 0.5), Ω) ≈ [z^-3/(1 + z^-5) for z in exp.(im*Ω*0.5)]
@test freqresp(c2d(feedback(0.5, delay(1.5)), 0.5), Ω) ≈ [0.5/(1 + 0.5*z^-3) for z in exp.(im*Ω*0.5)]

# Random conversions
sys1 = DelayLtiSystem(1.0/s)
@test sys1.P.A == sys1.P.D == fill(0,1,1)
@test sys1.P.B*sys1.P.C == fill(1,1,1)
@test sys1.Tau == []

## Multiple Delays
G = ss(1.0) + delay(2) + delay(3)
G_fr = 1 .+ exp.(-2*im*ω) .+ exp.(-3*im*ω)
@test freqresp(G, ω)[:] ≈ G_fr rtol=1e-15
@test freqresp(feedback(1.0, G), ω)[:] ≈ 1 ./(1 .+ G_fr) # Somewhat pathological system though

sys = feedback(1.0, P1 * (delay(2) + delay(3)))
expected_sys_fr = 1.0 ./ (P1_fr .* (exp.(-2*im*ω) + exp.(-3*im*ω)) .+ 1)

@test freqresp(sys, ω)[:] ≈ expected_sys_fr rtol=1e-14

@test freqresp(feedback(1.0*P1, G*G*P2*P2), ω)[:] ≈ P1_fr ./(1 .+ (G_fr.^2).*P1_fr.*P2_fr.^2) rtol=1e-15# Somewhat pathological system though

s11 = feedback(ss(1/s), delay(1))
s12 = ss(1/s)
s21 = DelayLtiSystem(ss(1/(s+1)))
s22 = ss(1/(s+10))

s1 = [s11 s12]
s2 = [s21 s22]

f1 = [s1;s2]
f2 = [s11 s12;
      s21 s22]

# Test that different consatenations work
w = 10 .^ (-2:0.1:2)
@test freqresp(f1, w) ≈ freqresp(f2, w) rtol=1e-15

# This used to be a weird bug
@test freqresp(s11, w) ≈ freqresp(f2[1,1], w) rtol=1e-15

@test propertynames(delay(1.0)) == (:P, :Tau)

#FIXME: A lot more tests, including MIMO systems in particular

# Test step
println("Simulating first delay system:")
@time step(delay(1)*tf(1,[1.,1]))
@time step(delay(1)*tf(1,[1,1]))

@time y1, t1, x1 = step([s11;s12], 10)
@time @test y1[2:2,:] ≈ step(s12, t1)[1] rtol = 1e-14

t = 0.0:0.1:10
y2, t2, x2 = step(s1, t)
# TODO Figure out which is inexact here
@test y2[1:1,:,1:1] + y2[1:1,:,2:2] ≈ step(s11, t)[1] + step(s12, t)[1] rtol=1e-14

y3, t3, x3 = step([s11; s12], t)
@test y3[1:1,:,1] ≈ step(s11, t)[1] rtol = 1e-14
@test y3[2:2,:,1] ≈ step(s12, t)[1] rtol = 1e-14

y1, t1, x1 = step(DelayLtiSystem([1.0/s 2/s; 3/s 4/s]), t)
y2, t2, x2 = step([1.0/s 2/s; 3/s 4/s], t)
@test y1 ≈ y2 rtol=1e-14
@test size(x1,2) == length(t)
@test size(x1,3) == 2

##################### Test with known solution
K = 2.0
sys_known = feedback(delay(1)*K/s, 1)

ystep, t, _ = step(sys_known, 3)

function y_expected(t, K)
      if t < 1
            return 0
      elseif t < 2
            return K*(t-1)
      elseif t <= 3
            return K*(t-1)-1/2*K^2*(t-2)^2
      else
            throw(ArgumentError("Test not defined here"))
      end
end

@test ystep' ≈ y_expected.(t, K) atol = 1e-12

function dy_expected(t, K)
      if t < 1
            return 0
      elseif t < 2
            return K
      elseif t <= 3
            return K - K^2*(t-2)
      else
            throw(ArgumentError("Test not defined here"))
      end
end

y_impulse, t, _ = impulse(sys_known, 3)

# TODO Better accuracy for impulse
@test y_impulse' ≈ dy_expected.(t, K) rtol=1e-13
@test maximum(abs, y_impulse' - dy_expected.(t, K)) < 1e-12

y_impulse, t, _ = impulse(sys_known, 3, alg=MethodOfSteps(Tsit5()))
# Two orders of magnitude better with BS3 in this case, which is default for impulse
@test y_impulse' ≈ dy_expected.(t, K) rtol=1e-5
@test maximum(abs, y_impulse' - dy_expected.(t, K)) < 1e-5

## Test delay with D22 term
t = 0:0.01:4

sys = delay(1)

y, t, x = step(sys, t)
@test y[:] ≈ [zeros(100); ones(301)] atol = 1e-12
@test size(x) == (0,401)

sys = delay(1)*delay(1)*1/s

y, t, x = step(sys, t)

y_sol = [zeros(200);0:0.01:2]'

@test maximum(abs,y-y_sol) < 1e-13
@test maximum(abs,x-collect(0:0.01:4)') < 1e-15

# TODO For some reason really bad accuracy here
# Looks like a lag in time
sys = 1/s*delay(1)*delay(1)

y, t, x = step(sys, t)
@test maximum(abs,y-y_sol) < 1e-5
@test maximum(abs,x-y_sol) < 1e-5

t = 0:0.001:0.1
y, t, x = step(sys, t)
@test length(y) == length(t)

##  Test of basic pade functionality

Ω = [0, 0.5, 1, 2, 5]
@test freqresp(pade(1, 1), Ω) == freqresp(tf([-1/2, 1], [1/2, 1]), Ω)
@test freqresp(pade(1, 2), Ω) ≈ freqresp(tf([1/12, -1/2, 1], [1/12, 1/2, 1]), Ω)

for (n, tol)=enumerate([0.05; 1e-3; 1e-5; 1e-7; 1e-11; 1e-14*ones(5)])
    G = pade(0.8, n)

    @test isstable(G)
    @test evalfr(G, 0)[1] ≈ 1
    @test abs(evalfr(G, 2im)[1]) ≈ 1
    @test evalfr(G, 1im)[1] ≈ exp(-0.8im) atol=tol
end

## Test pade applied to DelayLtiSystem

@test freqresp(pade(delay(0.5), 2), Ω) ≈ freqresp(pade(0.5, 2), Ω)

P = delay(1.0) * DemoSystems.lag(T=1)
@test freqresp(pade(feedback(1,P), 2), Ω) == freqresp(feedback(1, pade(P,2)), Ω)

Ω = [0, 0.1, 0.2]
P_wb = DemoSystems.woodberry()

@test freqresp(pade(P_wb, 2), Ω) ≈ freqresp(P_wb, Ω) atol=0.02
@test freqresp(pade(P_wb, 3), Ω) ≈ freqresp(P_wb, Ω) atol=5e-4

@test freqresp(pade(feedback(eye_(2), P_wb), 3), Ω) ≈ freqresp(feedback(eye_(2), P_wb), Ω) atol=1e-4

end


"""
    find_biclique_cover(skeleton, faces)

Given a skeleton of a planar graph and its faces, finds a biclique cover for the complement
of the respective finite element graph with a divide-and-conquer approach: construct a biclique
from a separator for the current finite element graph and call the function recursively on the
induced subskeletons and subfaces of A ∪ C and B ∪ C.
"""
function find_biclique_cover(skeleton::LabeledGraph{T}, faces::Set{Vector{T}})::Set{Pair{Set{T}, Set{T}}} where {T}
    face_pairs = _find_face_pairs(faces)
    feg = ClutteredEnvPathOpt._find_finite_element_graph(skeleton, face_pairs)

    if ClutteredEnvPathOpt.lg_is_complete(feg)
        return Set{Pair{Set{T}, Set{T}}}()
    end

    (C, A, B) = _find_feg_separator_lt_no_empty(skeleton, face_pairs)
    # println("Set A: $A\nSet B: $B\nSet C: $C")

    if !isempty(A) && !isempty(B)
        node = Set([A => B])
    else
        node = Set{Pair{Set{T}, Set{T}}}()
        println("A or B are empty.\nSkeleton vertices are $(skeleton.labels)")
    end

    if !isempty(union(A,C))
        skeleton_ac, faces_ac = _find_skeleton_faces(union(A, C), skeleton, faces)
        if !isempty(faces_ac)
            left = find_biclique_cover(skeleton_ac, faces_ac)
        else
            left = Set{Pair{Set{T}, Set{T}}}()
        end
    else
        left = Set{Pair{Set{T}, Set{T}}}()
    end

    if !isempty(union(B,C))
        skeleton_bc, faces_bc = _find_skeleton_faces(union(B, C), skeleton, faces)
        if !isempty(faces_bc)
            right = find_biclique_cover(skeleton_bc, faces_bc)
        else
            right = Set{Pair{Set{T}, Set{T}}}()
        end
    else
        right = Set{Pair{Set{T}, Set{T}}}()
    end

    return union(node, left, right)
end

"""
    find_biclique_cover_debug(skeleton, faces, points, obstacles, file_name)

Same as find_biclique_cover, with the added features of plotting the new skeletons and
recording separator information at each call.
"""
function find_biclique_cover_debug(skeleton::LabeledGraph{T}, faces::Set{Vector{T}}, points::Vector{Any}, obstacles, file_name)::Set{Pair{Set{T}, Set{T}}} where {T}
    ClutteredEnvPathOpt.plot_edges(skeleton, points, obstacles)
    f = open(file_name, "a")
    for (i,face) in enumerate(faces)
        write(f, "Face $i: $face\n")
        # println("Face $i: $face")
    end
    flush(f)
    close(f)

    face_pairs = _find_face_pairs(faces)
    feg = ClutteredEnvPathOpt._find_finite_element_graph(skeleton, face_pairs)

    if ClutteredEnvPathOpt.lg_is_complete(feg)
        f = open(file_name, "a")
        write(f, "FEG is complete.\n")
        flush(f)
        close(f)

        # println("FEG is complete.")
        return Set{Pair{Set{T}, Set{T}}}()
    end

    (C, A, B) = _find_feg_separator_lt_no_empty(skeleton, face_pairs)

    if !isempty(A) && !isempty(B)
        node = Set([A => B])
    else
        node = Set{Pair{Set{T}, Set{T}}}()
        f = open(file_name, "a")
        write(f, "A or B is empty.\n")
        flush(f)
        close(f)

        # println("A or B is empty.")
    end

    if !isempty(union(A,C))
        skeleton_ac, faces_ac = _find_skeleton_faces(union(A, C), skeleton, faces)
        if !isempty(faces_ac)  # I think this will always be nonempty of A ∪ C is nonempty
            f = open(file_name, "a")
            write(f, "\nPlot \nSet A: $A\nSet B: $B\nSet C: $C\nA ∪ C\n")
            flush(f)
            close(f)

            # println("\nSet A: $A\nSet B: $B\nSet C: $C\nA ∪ C")
            left = find_biclique_cover_debug(skeleton_ac, faces_ac, points, obstacles, file_name)
        else
            left = Set{Pair{Set{T}, Set{T}}}()
            f = open(file_name, "a")
            write(f, "faces_ac is empty.\n")
            flush(f)
            close(f)

            # println("faces_ac is empty.")
        end
    else
        left = Set{Pair{Set{T}, Set{T}}}()
        f = open(file_name, "a")
        write(f, "A ∪ C is empty.\n")
        flush(f)
        close(f)

        # println("A ∪ C is empty.")
    end

    if !isempty(union(B,C))
        skeleton_bc, faces_bc = _find_skeleton_faces(union(B, C), skeleton, faces)
        if !isempty(faces_bc)
            f = open(file_name, "a")
            write(f, "\nPlot \nSet A: $A\nSet B: $B\nSet C: $C\nB ∪ C\n")
            flush(f)
            close(f)

            # println("\nSet A: $A\nSet B: $B\nSet C: $C\nB ∪ C")
            right = find_biclique_cover_debug(skeleton_bc, faces_bc, points, obstacles, file_name)
        else
            right = Set{Pair{Set{T}, Set{T}}}()
            f = open(file_name, "a")
            write(f, "faces_bc is empty.\n")
            flush(f)
            close(f)

            # println("faces_bc is empty.")
        end
    else
        right = Set{Pair{Set{T}, Set{T}}}()
        f = open(file_name, "a")
        write(f, "B ∪ C is empty.\n")
        flush(f)
        close(f)

        # println("B ∪ C is empty.")
    end

    return union(node, left, right)
end

function find_biclique_cover_one_iter(skeleton::LabeledGraph{T}, faces::Set{Vector{T}}) where {T}
    face_pairs = _find_face_pairs(faces)
    feg = ClutteredEnvPathOpt._find_finite_element_graph(skeleton, face_pairs)

    if ClutteredEnvPathOpt.lg_is_complete(feg)
        return Set{Pair{Set{T}, Set{T}}}()
    end

    (C, A, B) = _find_feg_separator_lt_no_empty(skeleton, face_pairs)
    # println("Set A: $A\nSet B: $B\nSet C: $C")

    # if !isempty(A) && !isempty(B)
    #     node = Set([A => B])
    # else
    #     node = Set{Pair{Set{T}, Set{T}}}()
    #     println("A or B are empty.\nSkeleton vertices are $(skeleton.labels)")
    # end

    if !isempty(union(A,C))
        skeleton_ac, faces_ac = _find_skeleton_faces(union(A, C), skeleton, faces)
    #     if !isempty(faces_ac)
    #         left = find_biclique_cover(skeleton_ac, faces_ac)
    #     else
    #         left = Set{Pair{Set{T}, Set{T}}}()
    #     end
    # else
    #     left = Set{Pair{Set{T}, Set{T}}}()
    end

    if !isempty(union(B,C))
        skeleton_bc, faces_bc = _find_skeleton_faces(union(B, C), skeleton, faces)
    #     if !isempty(faces_bc)
    #         right = find_biclique_cover(skeleton_bc, faces_bc)
    #     else
    #         right = Set{Pair{Set{T}, Set{T}}}()
    #     end
    # else
    #     right = Set{Pair{Set{T}, Set{T}}}()
    end

    # return union(node, left, right)
    return (C,A,B), skeleton_ac, faces_ac, skeleton_bc, faces_bc
end

"""
    biclique_merger(cover, feg)

Given a biclique cover, merges bicliques together when possible for a smaller cover.
"""
function biclique_merger(cover::Set{Pair{Set{T}, Set{T}}}, feg::LabeledGraph{T})::Set{Pair{Set{T},Set{T}}} where {T}
    # Sort bicliques by size so that we attempt merging smaller bicliques first
    cover = sort(collect(cover), by=(s->length(s.first)+length(s.second)))

    # Loop until we are unable to find any merges
    merge_found = true
    while merge_found
        for i = 1:length(cover)
            for j = i+1:length(cover)
                # TODO: Check if this is right
                # if (!isempty(intersect(cover[i].first, cover[j].first)) && !isempty(intersect(cover[i].first, cover[j].second))) ||
                #     (!isempty(intersect(cover[i].second, cover[j].first)) && !isempty(intersect(cover[i].second, cover[j].second)))
                #     continue
                # end

                potential_merge_found = false
                merge_found = false
                A1_adj_A2 = false
                A1_adj_B2 = false
                B1_adj_A2 = false
                B1_adj_B2 = false
                for u in cover[i].first
                    for v in cover[j].first
                        if LightGraphs.has_edge(feg.graph, feg.labels[u], feg.labels[v])
                            A1_adj_A2 = true
                            break
                        end
                    end
                    if A1_adj_A2
                        break
                    end
                end

                for u in cover[i].first
                    for v in cover[j].second
                        if LightGraphs.has_edge(feg.graph, feg.labels[u], feg.labels[v])
                            A1_adj_B2 = true
                            break
                        end
                    end
                    if A1_adj_B2
                        break
                    end
                end

                if A1_adj_A2 && A1_adj_B2
                    continue  # cannot merge the bicliques in this case
                end

                if !A1_adj_A2
                    for u in cover[i].second
                        for v in cover[j].second
                            if LightGraphs.has_edge(feg.graph, feg.labels[u], feg.labels[v])
                                B1_adj_B2 = true
                                break
                            end
                        end
                        if B1_adj_B2
                            break
                        end
                    end
                end

                if !A1_adj_B2
                    for u in cover[i].second
                        for v in cover[j].first
                            if LightGraphs.has_edge(feg.graph, feg.labels[u], feg.labels[v])
                                B1_adj_A2 = true
                                break
                            end
                        end
                        if B1_adj_A2
                            break
                        end
                    end
                end

                # Two possible merges. If both are possible, will need to see if it matters which merge we choose
                A = Set([])
                B = Set([])
                if !A1_adj_A2 && !B1_adj_B2
                    # A = A1 ∪ B2
                    A = union(cover[i].first, cover[j].second)
                    # B = B1 ∪ A2
                    B = union(cover[i].second, cover[j].first)

                    potential_merge_found = true
                elseif !A1_adj_B2 && !B1_adj_A2
                    # A = A1 ∪ A2
                    A = union(cover[i].first, cover[j].first)
                    # B = B1 ∪ B2
                    B = union(cover[i].second, cover[j].second)

                    potential_merge_found = true
                end

                # Second condition is to prevent vertices from being in both A and B
                if potential_merge_found && isempty(intersect(A,B)) # don't need second condition if above is correct
                    merge_found = true
                    deleteat!(cover, j)
                    deleteat!(cover, i)
                    push!(cover, A => B)
                    sort!(cover, by=(s->length(s.first)+length(s.second)))
                    break
                end

            end
            if merge_found
                break
            end
        end

    end

    return Set(cover)
end

struct BicliqueCoverTree
    biclique
    left
    right
end

function find_biclique_cover_as_tree(skeleton::LabeledGraph{T}, faces::Set{Vector{T}}) where {T}
    face_pairs = ClutteredEnvPathOpt._find_face_pairs(faces)
    feg = ClutteredEnvPathOpt._find_finite_element_graph(skeleton, face_pairs)

    if ClutteredEnvPathOpt.lg_is_complete(feg)
        return nothing
    end

    (C, A, B) = ClutteredEnvPathOpt._find_feg_separator_lt_no_empty(skeleton, face_pairs)

    if !isempty(A) && !isempty(B)
        node = A => B
    else
        node = nothing
    end
    if !isempty(union(A,C))
        skeleton_ac, faces_ac = ClutteredEnvPathOpt._find_skeleton_faces(union(A, C), skeleton, faces)
        if !isempty(faces_ac)
            left = ClutteredEnvPathOpt.find_biclique_cover_as_tree(skeleton_ac, faces_ac)
        else
            left = nothing
        end
    else
        left = nothing
    end
    if !isempty(union(B,C))
        skeleton_bc, faces_bc = ClutteredEnvPathOpt._find_skeleton_faces(union(B, C), skeleton, faces)
        if !isempty(faces_bc)
            right = ClutteredEnvPathOpt.find_biclique_cover_as_tree(skeleton_bc, faces_bc)
        else
            right = nothing
        end
    else
        right = nothing
    end

    return BicliqueCoverTree(node, left, right)
end

function flatten(tree)
    if (tree === nothing)
        return []
    end

    return union([tree.biclique], flatten(tree.left), flatten(tree.right))
end

function getEdges(tree, cover)
    res = ""

    if (tree === nothing)
        return res
    end

    node_num = findfirst(x -> x === tree.biclique, cover)

    if (tree.left !== nothing)
        left_num = findfirst(x -> x === tree.left.biclique, cover)
        res = string(res, node_num, " -> ", left_num, " ", getEdges(tree.left, cover))
    end

    if (tree.right !== nothing)
        right_num = findfirst(x -> x === tree.right.biclique, cover)
        res = string(res, node_num, " -> ", right_num, " ", getEdges(tree.right, cover))
    end

    return res
end

function tree2digraph(tree)
    cover = collect(flatten(tree))

    edges = getEdges(tree, cover)

    return (cover, string("Digraph G { ", edges, " }"))
end

"""
    _find_skeleton_faces(vertices, old_skeleton, old_faces)

Find the skeleton and faces of a finite element subgraphgraph given a subset of
vertices, a finite element graph's skeleton and the faces of its planar
embedding. Returns a (skeleton, face vector) tuple.

Note: ensure the faces are passed in as a set of clockwise-ordered vertices
"""
function _find_skeleton_faces(vertices::Set{T}, old_skeleton::LabeledGraph{T}, old_faces::Set{Vector{T}})::Tuple{LabeledGraph{T}, Set{Vector{T}}} where {T}
    skeleton = copy(old_skeleton)
    for vertex in keys(old_skeleton.labels)
        if !in(vertex, vertices)
            ClutteredEnvPathOpt.lg_rem_vertex!(skeleton, vertex)
        end
    end

    # Remove obsolete vertices from each face and return faces with at least 1 vertex
    faces = @pipe map(face -> filter(vertex -> in(vertex, vertices), face), collect(old_faces)) |>
                filter(face -> length(face) > 0, _)

    # If a new face is a subset of another new face, don't include it
    new_faces = []
    removed_indices = []
    for i = 1:length(faces)
        indicator = 0
        for j = 1:length(faces)
            if j == i || j in removed_indices
                continue
            end
            if intersect(Set(faces[i]), Set(faces[j])) == Set(faces[i]) && length(faces[j]) >= length(faces[i])
                indicator += 1
                break
            end
        end
        if indicator != 0
            push!(removed_indices, i)
        else
            push!(new_faces, faces[i])
        end
    end

    # Close the new subface
    for face in new_faces
        if length(face) > 1
            for i = 2:length(face)
                ClutteredEnvPathOpt.lg_add_edge!(skeleton, face[i-1], face[i])
            end
            ClutteredEnvPathOpt.lg_add_edge!(skeleton, face[end], face[1])
        end
    end

    return skeleton, Set(new_faces)
end

"""
    _find_face_pairs(faces)

Convert a set of list of vertices included in a face to a set of sets of edges
comprising a face.
"""
function _find_face_pairs(faces::Set{Vector{T}})::Set{Set{Pair{T, T}}} where {T}
    face_pairs = Set{Set{Pair{T, T}}}()
    for face in faces
        pairs = Set{Pair{T, T}}()

        for i in 2:length(face)
            push!(pairs, face[i - 1] => face[i])
        end

        push!(pairs, face[end] => face[1])

        push!(face_pairs, pairs)
    end

    return face_pairs
end

"""
    _find_feg_separator_lt_no_empty(skeleton, face_pairs)

Find the separator of a finite element graph, repeating the process if for some
given root either A or B is empty. Returns the separator for which |A|*|B| is largest.
"""
function _find_feg_separator_lt_no_empty(skeleton::LabeledGraph{T}, face_pairs::Set{Set{Pair{T, T}}})::Tuple{Set{T}, Set{T}, Set{T}} where {T}
    feg = ClutteredEnvPathOpt._find_finite_element_graph(skeleton, face_pairs)
    # for root in keys(skeleton.labels)
    #     (C, A, B) = find_feg_separator_lt(skeleton, face_pairs, root)
    #     (C, A, B) = pp_expell(feg, C, A, B)

    #     if !isempty(A) && !isempty(B)
    #         return (C, A, B)
    #     end
    # end

    # return (Set([]), Set([]), Set([]))
    #####################################################################################################

    separators = @pipe map(root -> find_feg_separator_lt(skeleton, face_pairs, root), collect(keys(skeleton.labels))) |>
        map(sep -> pp_expell(feg, sep[1], sep[2], sep[3]), _)

    balances = map(separator -> begin
        if isempty(separator[2]) || isempty(separator[3]) # A or B empty is no bueno
            return -1.0
        elseif isempty(separator[1]) # C empty, A and B are not, prioritize this case
            return 1.0 + min(length(separator[2]) / length(separator[3]), length(separator[3]) / length(separator[2]))
        else # None are empty
            return 1.0  # the separator that finds the most edges will be chosen
            # return min(length(separator[2]) / length(separator[3]), length(separator[3]) / length(separator[2]))
        end
    end, separators)

    max_balance = max(balances...)
    # println("Max Balance: $max_balance")

    if max_balance < -0.5
        # If at least one of A or B is empty, the subskeleton would be the same and we'll have infinite recursion
        # In such a case, our SEPARATOR ALGORITHM and PP_EXPELL cannot produce two nonempty A and B
        # Return empty sets since there'll be no biclique
        return (Set([]), Set([]), Set([]))
    else
        # Return the biclique that makes the most edges (e.g. |A|*|B| is largest)
        select = filter(b -> balances[b] == max_balance, 1:length(balances))
        max_edges = 0
        best_index = 0
        for i in select
            if length(separators[i][2])*length(separators[i][3]) > max_edges
                max_edges = length(separators[i][2])*length(separators[i][3])
                best_index = i
            end
        end
        return separators[best_index]
    end

end

"""
    _is_valid_biclique_cover(lg, cover)

Tests whether or not a cover is a valid biclique cover of a given graph.
"""
function _is_valid_biclique_cover(lg::LabeledGraph{T}, cover::Set{Pair{Set{T}, Set{T}}})::Bool where {T}
    e_bar = LightGraphs.edges(LightGraphs.complement(lg.graph))

    edges = @pipe map(pair -> _cartesian_product(pair.first, pair.second), collect(cover)) |> reduce(union!, _, init=Set{Pair{T, T}}())
    test = LightGraphs.SimpleGraph(LightGraphs.nv(lg.graph))

    for edge in edges
        # edge.first/second will be our nodes, so we map them to the nodes they are in the graph
        LightGraphs.add_edge!(test, lg.labels[edge.first], lg.labels[edge.second])
    end

    return LightGraphs.edges(test) == e_bar
end

"""
    _is_valid_biclique_cover_diff(lg, cover)

Tests whether or not a cover is a valid biclique cover of a given graph. Returns
their difference in cardinality (positive => missing edges, negative => surplus of edges)
and a LabeledGraph containing any missing edges.
"""
function _is_valid_biclique_cover_diff(lg::LabeledGraph{T}, cover::Set{Pair{Set{T}, Set{T}}})::Tuple{Bool,Number,Set{LightGraphs.SimpleGraphs.SimpleEdge{T}},LabeledGraph{T}} where {T}
    e_bar = LightGraphs.edges(LightGraphs.complement(lg.graph))

    edges = @pipe map(pair -> _cartesian_product(pair.first, pair.second), collect(cover)) |> reduce(union!, _, init=Set{Pair{T, T}}())
    test = LightGraphs.SimpleGraph(LightGraphs.nv(lg.graph))
    rev = _reverse_labels(lg.labels)

    for edge in edges
        # edge.first/second will be our nodes, so we map them to the nodes they are in the graph
        LightGraphs.add_edge!(test, lg.labels[edge.first], lg.labels[edge.second])
    end

    # Create a LabeledGraph containing only the missing edges
    set_diff = setdiff(Set(e_bar), Set(LightGraphs.edges(test)))
    missing_e_graph = LightGraphs.SimpleGraph(LightGraphs.nv(lg.graph))
    missing_e_lg = LabeledGraph{Int}(missing_e_graph, lg.labels)
    for edge in set_diff
        lg_add_edge!(missing_e_lg, rev[edge.src], rev[edge.dst])
    end

    return LightGraphs.edges(test) == e_bar, length(e_bar) - LightGraphs.ne(test), set_diff, missing_e_lg
end

"""
    _cartesian_product(A, B)

Returns the cartesian product of two sets
"""
function _cartesian_product(A::Set{V}, B::Set{W})::Set{Pair{V, W}} where {V, W}
    res = Set{Pair{V, W}}()

    for a in A
        for b in B
            push!(res, Pair(a, b))
        end
    end

    return res
end

"""
    _find_feg_separator_lt_no_empty_visual(skeleton, face_pairs, locations)
Finds (C, A, B) and plots the edges found between A and B.
"""
function _find_feg_separator_lt_no_empty_visual(skeleton::LabeledGraph{T}, face_pairs::Set{Set{Pair{T, T}}}, points::Vector{Any})::Tuple{Set{T}, Set{T}, Set{T}} where {T}
    (C,A,B) = _find_feg_separator_lt_no_empty(skeleton, face_pairs)

    if isempty(A) || isempty(B)
        return (C,A,B)
    end

    # Create edges between A and B
    new_edges = ClutteredEnvPathOpt._cartesian_product(A, B)

    Plots.plot()
    Plots.scatter!(map(point -> point.first, points), map(point -> point.second, points))

    for a in A
        Plots.scatter!([points[a].first], [points[a].second], color=:red)
    end
    for b in B
        Plots.scatter!([points[b].first], [points[b].second], color=:yellow)
    end

    for edge in new_edges
        (x1,y1) = points[edge.first]
        (x2,y2) = points[edge.second]
        Plots.plot!([x1, x2], [y1, y2], color=:green, linestyle=:dash)
    end

    display(Plots.plot!(legend=:false,xlims=(-.05,1.05),ylims=(-0.05,1.05), title="Biclique"))

    return (C,A,B)
end

"""
    _find_biclique_cover_visual(skeleton, faces, locations)
Finds a biclique cover and plots the edges of each biclique.
"""
function _find_biclique_cover_visual(skeleton::LabeledGraph{T}, faces::Set{Vector{T}}, points::Vector{Any})::Set{Pair{Set{T}, Set{T}}} where {T}
    face_pairs = _find_face_pairs(faces)
    feg = ClutteredEnvPathOpt._find_finite_element_graph(skeleton, face_pairs)

    if ClutteredEnvPathOpt.lg_is_complete(feg)
        return Set{Pair{Set{T}, Set{T}}}()
    end

    (C, A, B) = _find_feg_separator_lt_no_empty_visual(skeleton, face_pairs, points)
    # println("Set A: $A\nSet B: $B\nSet C: $C")

    if !isempty(A) && !isempty(B)
        node = Set([A => B])
    else
        node = Set{Pair{Set{T}, Set{T}}}()
    end
    if !isempty(union(A,C))
        skeleton_ac, faces_ac = _find_skeleton_faces(union(A, C), skeleton, faces)
        if !isempty(faces_ac)
            # println("\nSet A: $A\nSet B: $B\nSet C: $C\nA ∪ C")
            left = _find_biclique_cover_visual(skeleton_ac, faces_ac, points)
        else
            left = Set{Pair{Set{T}, Set{T}}}()
        end
    else
        left = Set{Pair{Set{T}, Set{T}}}()
    end
    if !isempty(union(B,C))
        skeleton_bc, faces_bc = _find_skeleton_faces(union(B, C), skeleton, faces)
        if !isempty(faces_bc)
            # println("\nSet A: $A\nSet B: $B\nSet C: $C\nB ∪ C")
            right = _find_biclique_cover_visual(skeleton_bc, faces_bc, points)
        else
            right = Set{Pair{Set{T}, Set{T}}}()
        end
    else
        right = Set{Pair{Set{T}, Set{T}}}()
    end

    return union(node, left, right)
end


@testset "Initial conditions" begin
    @testset "Fluid Sine Wave" begin
        h = ones(10,10)
        Swalbe.sinewave1D!(h, 2.0, 1, 1, 1)
        @test maximum(h) ≈ 2*(1 + 1) atol= 0.5
        @test minimum(h) ≈ 2*(1 - 1) atol= 0.5
    end

    @testset "Random interface" begin
        h = ones(10,10)
        Swalbe.randinterface!(h, 2.0, 0.1)
        @test maximum(h) <= 3.0
        @test minimum(h) >= 1.0
    end

    @testset "Single droplet" begin
        rad = 50
        θ = 1/3
        center = (50,50)
        height = Swalbe.singledroplet(ones(100,100), rad, θ, center)
        @test isa(height, Array{Float64,2})
        @test size(height) == (100, 100)
        @test findmax(height)[1] == rad * (1 - cospi(θ))
        @test findmax(height)[2] == CartesianIndex(center[1],center[2])
        # 1D case
        rad = 50
        θ = 1/3
        center = 100
        height = Swalbe.singledroplet(ones(200), rad, θ, center)
        @test isa(height, Vector{Float64})
        @test size(height) == (200,)
        @test findmax(height)[1] == rad * (1 - cospi(θ))
        @test findmax(height)[2] == center
    end

    @testset "Restart height" begin
        h1 = rand(10,10)
        h2 = rand(10,10)
        f = zeros(2, 100)
        f = Dict()
        f["h_1"] = vec(h1)
        f["h_2"] = vec(h2)
        bson("file.bson", f)
        save("file.jld2", f)
        h = Swalbe.restart_from_height("file.bson", kind="bson", timestep=0, size=(10,10))
        @test all(h2 .== h)
        h = Swalbe.restart_from_height("file.bson", kind="bson", timestep=1, size=(10,10))
        @test all(h1 .== h)
        h = Swalbe.restart_from_height("file.jld2", size=(10,10))
        @test all(h2 .== h)
        h = Swalbe.restart_from_height("file.jld2", kind="jld2", timestep=1, size=(10,10))
        @test all(h1 .== h)
        rm("file.bson")
        rm("file.jld2")
    end
end
@testset "Contact angle patterns" begin
    @testset "Default elliptical patch" begin
        theta, P = Swalbe.ellipsepattern(ones(100,100), 1/9) # with kw: center=(Int(lx/2), Int(ly/2)),δₐ=1/36, a=10, b=5
        @test size(theta) == (100, 100)
        @test theta[50,50] == 1/9 + 1/36
        @test theta[1,1] == 1/9
    end
    @testset "Off center elliptical patch" begin
        theta, P = Swalbe.ellipsepattern(ones(100,100), 1/9, center=(20,20)) # with kw: center=(Int(lx/2), Int(ly/2)),δₐ=1/36, a=10, b=5
        @test size(theta) == (100, 100)
        @test theta[20,20] == 1/9 + 1/36
        @test theta[100,100] == 1/9
    end
    @testset "Negative contrast elliptical patch" begin
        theta, P = Swalbe.ellipsepattern(ones(100,100), 1/9, δₐ=-1/36) # with kw: center=(Int(lx/2), Int(ly/2)),δₐ=1/36, a=10, b=5
        @test size(theta) == (100, 100)
        @test theta[50,50] == 1/9 - 1/36
        @test theta[100,100] == 1/9
    end
    @testset "Larger elliptical patch" begin
        theta, P = Swalbe.ellipsepattern(ones(100,100), 1/9, a=10, b=20) # with kw: center=(Int(lx/2), Int(ly/2)),δₐ=1/36, a=10, b=5
        @test size(theta) == (100, 100)
        @test theta[40,50] == 1/9 + 1/36
        @test theta[70,50] == 1/9 + 1/36
        @test theta[100,100] == 1/9
    end
    @testset "Polygon patch" begin
        lx, ly = 100, 100
        @testset "Equilateral Triangle" begin
            θₙ, P = Swalbe.trianglepattern(ones(100, 100), 1/9)
            @test all(θₙ .> 0.0)
            @test θₙ[1,1] == 1/9 # the default value
            @test θₙ[Int(lx/2), Int(ly/2)] == 1/9 + 1/36
            θₙ, P = Swalbe.trianglepattern(ones(100, 100), 1/9, center=(lx/3, 2*ly/3))
            @test θₙ[1,1] == 1/9 # the default value
            @test θₙ[Int(round(lx/3)), Int(round(2*ly/3))] == 1/9 # Because the side became to large
            θₙ, P = Swalbe.trianglepattern(ones(100, 100), 1/9, center=(0, 2*ly/3))
            @test θₙ[1,1] == 1/9 # the default value
            @test θₙ[Int(round(lx/3)), Int(round(2*ly/3))] == 1/9 # Because the side became to large
            @test θₙ[60, 1] == 1/9 + 1/36 # If too large stick to origin!
            θₙ, P = Swalbe.trianglepattern(ones(100, 100), 1/9, δₐ = -1/36)
            @test θₙ[1,1] == 1/9 # the default value
            @test θₙ[Int(lx/2), Int(ly/2)] == 1/9 - 1/36
            θₙ, P = Swalbe.trianglepattern(ones(100, 100), 1/9, side = 30) # Smaller triangle
            @test θₙ[1,1] == 1/9 # the default value
            @test θₙ[Int(lx/2), Int(ly/2)] == 1/9 + 1/36
            @test θₙ[50, 60] == 1/9 + 1/36
        end

        @testset "Box pattern" begin
            θₙ, P = Swalbe.boxpattern(ones(100, 100), 1/9)
            @test all(θₙ .> 0.0)
            @test θₙ[1,1] == 1/9 # the default value
            @test θₙ[Int(lx/2), Int(ly/2)] == 1/9 + 1/36
            θₙ, P = Swalbe.boxpattern(ones(100, 100), 1/9, center=(lx/3, 2*ly/3))
            @test θₙ[1,1] == 1/9 # the default value
            @test θₙ[Int(round(2*ly/3)), Int(round(lx/3))] == 1/9 + 1/36
            θₙ, P = Swalbe.boxpattern(ones(100, 100), 1/9, center=(0, 2*ly/3))
            @test θₙ[20,1] == 1/9 # the default value
            @test θₙ[60,1] == 1/9 + 1/36
            θₙ, P = Swalbe.boxpattern(ones(100, 100), 1/9, δₐ = -1/36)
            @test θₙ[1,1] == 1/9 # the default value
            @test θₙ[Int(lx/2), Int(ly/2)] == 1/9 - 1/36
            θₙ, P = Swalbe.boxpattern(ones(100, 100), 1/9, side = 30) # Smaller triangle
            @test θₙ[1,1] == 1/9 # the default value
            @test θₙ[Int(lx/2), Int(ly/2)] == 1/9 + 1/36
            @test θₙ[Int(lx/2) + Int(30/2), Int(ly/2) + Int(30/2)] == 1/9 + 1/36
        end
    end
end


using JSON
using Colors
using Gadfly
using DataFrames

using Mamba
using Stan

using PSIS

include("cvit.jl")

# Data
data = JSON.parsefile("wells.data.json")
y = Float64.(data["switched"])
x = Float64[data["arsenic"]  data["dist"]]
n, m = size(x)

# Model
model_str = readstring(open("arsenic_logistic.stan"))
stanmodel = Stanmodel(name="arsenic_logistic", adapt=500, update=500, model=model_str)

if isfile("sim.jls")
    sim = read("sim.jls", Chains)
else
    standata = [Dict("p" => m, "N" => n, "y" => y, "x" => x)]
    # Fit the model in Stan
    sim = stan(stanmodel, standata, '.', CmdStanDir=CMDSTAN_HOME, summary=false)
    write("sim.jls", sim)
end

names_sim = filter(x->startswith(x,"log_lik"), sim.names)

# Compute LOO and standard error
log_lik = sim[:, names_sim, :]
loo, loos, pk = psisloo(log_lik)
elpd_loo = sum(loos)
se_elpd_loo = std(loos) * sqrt(n)
@printf(">> elpd_loo = %.1f, SE(elpd_loo) = %.1f\n", elpd_loo, se_elpd_loo)

# Check the shape parameter k of the generalized Pareto distribution
if all(pk .< 0.5)
    println("All Pareto k estimates OK (k < 0.5)")
else
    pkn1 = sum((pk .>= 0.5) & (pk .< 1))
    pkn2 = sum(pk .>= 1)
    @printf(">> %d (%.0f%%) PSIS Pareto k estimates between 0.5 and 1\n", pkn1, pkn1/n*100)
    @printf(">> %d (%.0f%%) PSIS Pareto k estimates greater than 1\n", pkn2, pkn2/n*100)
end

exit()
# Fit a second model, using log(arsenic) instead of arsenic
x2 = Float64[log.(data["arsenic"])  data["dist"]]

# Model
if isfile("sim2.jls")
    sim2 = read("sim2.jls", Chains)
else
    standata2 = [Dict("p" => m, "N" => n, "y" => y, "x" => x2)]
    # Fit the model in Stan
    sim2 = stan(stanmodel, standata2, '.', CmdStanDir=CMDSTAN_HOME, summary=false)
    write("sim2.jls", sim2)
end

# Compute LOO and standard error
log_lik = sim2[:, names_sim, :]
loo2, loos2, pk2 = psisloo(log_lik)
elpd_loo = sum(loos2)
se_elpd_loo = std(loos2) * sqrt(n)
@printf(">> elpd_loo = %.1f, SE(elpd_loo) = %.1f\n", elpd_loo, se_elpd_loo)

# Check the shape parameter k of the generalized Pareto distribution
if all(pk .< 0.5)
    println("All Pareto k estimates OK (k < 0.5)")
else
    pkn1 = sum((pk .>= 0.5) & (pk .< 1))
    pkn2 = sum(pk .>= 1)
    @printf(">> %d (%.0f%%) PSIS Pareto k estimates between 0.5 and 1\n", pkn1, pkn1/n*100)
    @printf(">> %d (%.0f%%) PSIS Pareto k estimates greater than 1\n", pkn2, pkn2/n*100)
end

## Compare the models
loodiff = loos - loos2
@printf("elpd_diff = %.1f, SE(elpd_diff) = %.1f\n",sum(loodiff), std(loodiff) * sqrt(n))

## k-fold-CV
# k-fold-CV should be used if several khats>0.5
# in this case it is not needed, but provided as an example
model_str = readstring(open("arsenic_logistic_t.stan"))
stanmodel = Stanmodel(name="arsenic_logistic_t", adapt=500, update=500, model=model_str);

cvitr, cvitst = cvit(n, 10, true)
kfcvs = similar(loos)
for cvi in 1:10
    @printf("%d\n", cvi)

    standatacv = [Dict("p" => m, "N" => length(cvitr[cvi]), "Nt" => length(cvitst[cvi]),
                        "x" => x[cvitr[cvi],:], "y" => y[cvitr[cvi]],
                        "xt" => x[cvitst[cvi],:], "yt" => y[cvitst[cvi]])
                 ]
    # Fit the model in Stan
    simcv = stan(stanmodel, standatacv, '.',
                 CmdStanDir=CMDSTAN_HOME, summary=false)
    ns = filter(x->startswith(x,"log_likt"), simcv.names)
    log_likt = Mamba.combine(simcv[:, ns, :])
    kfcvs[cvitst[cvi]]= PSIS.logsumexp(log_likt) - log(size(log_likt,1))
end

# compare PSIS-LOO and k-fold-CV
p = plot(layer(x = loos, y = kfcvs, Geom.point),
     layer(x = [-3.5,0] ,y=[-3.5,0], Geom.line, style(default_color=colorant"red")),
     Guide.xlabel("PSIS-LOO"),
     Guide.ylabel("10-fold-CV"))

draw(PDF("Compare.pdf", 210mm, 210mm),p)


# # RealNVP network
# For the definition of this network and concepts of normalizing flow,
# please refer this realnvp blog: https://lilianweng.github.io/lil-log/2018/10/13/flow-based-deep-generative-models.html,
# and the pytorch notebook: https://github.com/GiggleLiu/marburg/blob/master/solutions/realnvp.ipynb

using NiLang, NiLang.AD
using LinearAlgebra
using DelimitedFiles
using Plots

# `include` the optimizer, you can find it under the `Adam.jl` file in the `examples/` folder.
include(NiLang.project_relative_path("examples", "Adam.jl"))

# ## Model definition
# First, define the single layer transformation and its behavior under `GVar` - the gradient wrapper.
struct RealNVPLayer{T}
    ## transform network
    W1::Matrix{T}
    b1::Vector{T}
    W2::Matrix{T}
    b2::Vector{T}
    y1::Vector{T}
    y1a::Vector{T}

    ## scaling network
    sW1::Matrix{T}
    sb1::Vector{T}
    sW2::Matrix{T}
    sb2::Vector{T}
    sy1::Vector{T}
    sy1a::Vector{T}
end

"""collect parameters in the `layer` into a vector `out`."""
function collect_params!(out, layer::RealNVPLayer)
    k=0
    for field in [:W1, :b1, :W2, :b2, :sW1, :sb1, :sW2, :sb2]
        v = getfield(layer, field)
        nv = length(v)
        out[k+1:k+nv] .= vec(v)
        k += nv
    end
    return out
end

"""dispatch vectorized parameters `out` into the `layer`."""
function dispatch_params!(layer::RealNVPLayer, out)
    k=0
    for field in [:W1, :b1, :W2, :b2, :sW1, :sb1, :sW2, :sb2]
        v = getfield(layer, field)
        nv = length(v)
        vec(v) .= out[k+1:k+nv]
        k += nv
    end
    return out
end

function nparameters(n::RealNVPLayer)
    sum(x->length(getfield(n, x)), [:W1, :b1, :W2, :b2, :sW1, :sb1, :sW2, :sb2])
end

# Then, we define `network` and how to access the parameters.
const RealNVP{T} = Vector{RealNVPLayer{T}}

nparameters(n::RealNVP) = sum(nparameters, n)

function collect_params(n::RealNVP{T}) where T
    out = zeros(T, nparameters(n))
    k = 0
    for layer in n
        np = nparameters(layer)
        collect_params!(view(out, k+1:k+np), layer)
        k += np
    end
    return out
end

function dispatch_params!(network::RealNVP, out)
    k = 0
    for layer in network
        np = nparameters(layer)
        dispatch_params!(layer, view(out, k+1:k+np))
        k += np
    end
    return network
end

function random_realnvp(nparams::Int, nhidden::Int, nhidden_s::Int, nlayer::Int; scale=0.1)
    random_realnvp(Float64, nparams, nhidden, nhidden_s::Int, nlayer; scale=scale)
end

function random_realnvp(::Type{T}, nparams::Int, nhidden::Int, nhidden_s::Int, nlayer::Int; scale=0.1) where T
    nin = nparams÷2
    scale = T(scale)
    y1 = zeros(T, nhidden)
    sy1 = zeros(T, nhidden_s)
    RealNVPLayer{T}[RealNVPLayer(
            randn(T, nhidden, nin)*scale, randn(T, nhidden)*scale,
            randn(T, nin, nhidden)*scale, randn(T, nin)*scale, y1, zero(y1),
            randn(T, nhidden_s, nin)*scale, randn(T, nhidden_s)*scale,
            randn(T, nin, nhidden_s)*scale, randn(T, nin)*scale, sy1, zero(sy1),
            ) for _ = 1:nlayer]
end

# ## Loss function
#
# In each layer, we use the information in `x` to update `y!`.
# During computing, we use to vector type ancillas `y1` and `y1a`,
# both of them can be uncomputed at the end of the function.

@i function onelayer!(x::AbstractVector{T}, layer::RealNVPLayer{T},
                y!::AbstractVector{T}, logjacobian!::T; islast) where T
    @routine @invcheckoff begin
        ## scale network
        scale ← zero(y!)
        ytemp2 ← zero(y!)
        i_affine!(layer.sy1, layer.sW1, layer.sb1, x)
        @inbounds for i=1:length(layer.sy1)
            if (layer.sy1[i] > 0, ~)
                layer.sy1a[i] += layer.sy1[i]
            end
        end
        i_affine!(scale, layer.sW2, layer.sb2, layer.sy1a)

        ## transform network
        i_affine!(layer.y1, layer.W1, layer.b1, x)
        ## relu
        @inbounds for i=1:length(layer.y1)
            if (layer.y1[i] > 0, ~)
                layer.y1a[i] += layer.y1[i]
            end
        end
    end
    ## inplace multiply exp of scale! -- dangerous
    @inbounds @invcheckoff for i=1:length(scale)
        @routine begin
            expscale ← zero(T)
            tanhscale ← zero(T)
            if (islast, ~)
                tanhscale += tanh(scale[i])
            else
                tanhscale += scale[i]
            end
            expscale += exp(tanhscale)
        end
        logjacobian! += tanhscale
        ## inplace multiply!!!
        temp ← zero(T)
        temp += y![i] * expscale
        SWAP(temp, y![i])
        temp -= y![i] / expscale
        temp → zero(T)
        ~@routine
    end

    ## affine the transform layer
    i_affine!(y!, layer.W2, layer.b2, layer.y1a)
    ~@routine
    ## clean up accumulated rounding error, since this memory is reused.
    @safe layer.y1 .= zero(T)
    @safe layer.sy1 .= zero(T)
end

# A realnvp network always transforms inputs reversibly.
# We update one half of `x!` a time, so that input and output memory space do not clash.
@i function realnvp!(x!::AbstractVector{T}, network::RealNVP{T}, logjacobian!) where T
    @invcheckoff for i=1:length(network)
        np ← length(x!)
        if (i%2==0, ~)
            @inbounds onelayer!(x! |> subarray(np÷2+1:np), network[i], x! |> subarray(1:np÷2), logjacobian!; islast=i==length(network))
        else
            @inbounds onelayer!(x! |> subarray(1:np÷2), network[i], x! |> subarray(np÷2+1:np), logjacobian!; islast=i==length(network))
        end
    end
end

# How to obtain the log-probability of a data.

@i function logp!(out!::T, x!::AbstractVector{T}, network::RealNVP{T}) where T
    (~realnvp!)(x!, network, out!)
    @invcheckoff for i = 1:length(x!)
        @routine begin
            xsq ← zero(T)
            @inbounds xsq += x![i]^2
        end
        out! -= 0.5 * xsq
        ~@routine
    end
end

# The negative-log-likelihood loss function

@i function nll_loss!(out!::T, cum!::T, xs!::Matrix{T}, network::RealNVP{T}) where T
    @invcheckoff for i=1:size(xs!, 2)
        @inbounds logp!(cum!, xs! |> subarray(:,i), network)
    end
    out! -= cum!/(@const size(xs!, 2))
end

# ## Training

function train(x_data, model; num_epochs = 800)
    num_vars = size(x_data, 1)
    params = collect_params(model)
    optimizer = Adam(; lr=0.01)
    for epoch = 1:num_epochs
        loss, a, b, c = nll_loss!(0.0, 0.0, copy(x_data), model)
        if epoch % 50 == 1
            println("epoch = $epoch, loss = $loss")
            display(showmodel(x_data, model))
        end
        _, _, _, gmodel = (~nll_loss!)(GVar(loss, 1.0), GVar(a), GVar(b), GVar(c))
        g = grad.(collect_params(gmodel))
        update!(params, grad.(collect_params(gmodel)), optimizer)
        dispatch_params!(model, params)
    end
    return model
end

function showmodel(x_data, model; nsamples=2000)
    scatter(x_data[1,1:nsamples], x_data[2,1:nsamples]; xlims=(-5,5), ylims=(-5,5))
    zs = randn(2, nsamples)
    for i=1:nsamples
        realnvp!(view(zs, :, i), model, 0.0)
    end
    scatter!(zs[1,:], zs[2,:])
end

# you can find the training data in `examples/` folder
x_data = Matrix(readdlm(NiLang.project_relative_path("examples", "train.dat"))')

import Random; Random.seed!(22)
model = random_realnvp(Float64, size(x_data, 1), 10, 10, 4; scale=0.1)

# Before training, the distribution looks like
# ![before](../../asset/nice_before.png)
model = train(x_data, model; num_epochs=800)

# After training, the distribution looks like
# ![before](../../asset/realnvp_after.png)


import JuMP: DiffableFunction

function frank_wolfe(p::CombiProblem{DiffableFunction})
  # check only constraint is one combinatorial constraint
  @assert length(p.convex_constraints==0) && length(p.combi_constraints==1)
  combiset = p.combi_constraints[1]
  n = combiset.dim
  x = zeros(n)
  for k=1:100 # maxiters
    g = grad(p.objective, x)
    h = fenchel(combiset, g)
    alpha = 1/k # stepsize
    x = (1-alpha)*x + alpha*h
  end
  return x
end


__precompile__()
module HierarchSingularIntegralEquations
    using Base, BandedMatrices, ApproxFun, SingularIntegralEquations, DualNumbers, RecipesBase,
            LinearAlgebra, Random, SpecialFunctions, LowRankApprox, InteractiveUtils

import Base: values, getindex, setindex!, *, +, -, ==, <, <=, >,
                >=, /, ^, \, ∪, transpose, convert, Array, Vector, Matrix

import Base.Broadcast: broadcasted, DefaultArrayStyle

import LinearAlgebra: ldiv!, mul!, rank, cond, qr

include("LinearAlgebra/LinearAlgebra.jl")
include("Operators/Operators.jl")
include("greensfun.jl")
include("fractals.jl")
include("clustertree.jl")

# if isdir(Pkg.dir("TikzGraphs"))
#     include("introspect.jl")
# end

include("plot.jl")
include("show.jl")

end #module