text:
using Logging
using Qaintessent
# using Qaintellect
include("/home/ga53vuw/Documents/PhD/projects/QAI/Qaintellect.jl/src/Qaintellect.jl")
using LinearAlgebra
using SparseArrays
using Flux
using IterTools: ncycle
# visualization
using Plots
using LaTeXStrings
include("./util.jl")
include("./adjacency_map.jl")
include("./thermal_states.jl")
include("./construct_hamiltonian.jl")

"""
Identity map as sparse matrix with real-valued entries.
"""
sparse_identity(n) = sparse(1.0*I, n, n)

global debug = false
global visualize = true

global Nx = 2
global Ny = 3
global isboundary = true

@assert Nx > 0 "Dimension of lattice must be positive and greater than 0"
@assert Ny > 0 "Dimension of lattice must be positive and greater than 0"

function single_qubit_layer(L::Integer)
    # local gates
    # [circuit_gate(i, RotationGate(0.05*π*randn(3))) for i in 1:L]
    cgc = CircuitGate[]
    append!(cgc, [circuit_gate(i, RxGate(0.05*π*randn())) for i in 1:L])
    append!(cgc, [circuit_gate(i, RyGate(0.05*π*randn())) for i in 1:L])
    append!(cgc, [circuit_gate(i, RzGate(0.05*π*randn())) for i in 1:L])
    cgc
end

"""
Construct a parametrized quantum layer consisting of single and two qubit rotation gates.
"""
function trotterized_layer_1d(L::Integer, istart::Integer; boundary::Bool=false)
    # interaction gates
    vcat(
        [circuit_gate(i, i+1, EntanglementXXGate(0.05π*randn())) for i in istart:2:L-1],
        ifelse(boundary, [circuit_gate(L, 1, EntanglementXXGate(0.05π*randn()))], []),
        [circuit_gate(i, i+1, EntanglementYYGate(0.05π*randn())) for i in istart:2:L-1],
        ifelse(boundary, [circuit_gate(L, 1, EntanglementYYGate(0.05π*randn()))], []),
        [circuit_gate(i, i+1, EntanglementZZGate(0.05π*randn())) for i in istart:2:L-1],
        ifelse(boundary, [circuit_gate(L, 1, EntanglementZZGate(0.05π*randn()))], [])
    )
end

function trotterized_layer_2d_h(Nx::Integer, Ny::Integer, istart::Integer; boundary::Bool=false)
    # interaction gates
    h = Iterators.flatten(repeat([i for i in istart:2:Nx-1], 1, Ny) .+ transpose([0:Ny-1...].*Nx))
    h_boundary = fill(Nx, Ny) .+ [0:Ny-1...].*Nx
    L = Nx * Ny
    vcat(
        [circuit_gate(i, i+1, EntanglementXXGate(0.05π*randn())) for i in h],
        ifelse(boundary&&(isodd(Nx)==isodd(istart)), [circuit_gate(i, i-Nx+1, EntanglementXXGate(0.05π*randn())) for i in h_boundary], []),
        [circuit_gate(i, i+1, EntanglementYYGate(0.05π*randn())) for i in h],
        ifelse(boundary&&(isodd(Nx)==isodd(istart)), [circuit_gate(i, i-Nx+1, EntanglementYYGate(0.05π*randn())) for i in h_boundary], []),
        [circuit_gate(i, i+1, EntanglementZZGate(0.05π*randn())) for i in h],
        ifelse(boundary&&(isodd(Nx)==isodd(istart)), [circuit_gate(i, i-Nx+1, EntanglementZZGate(0.05π*randn())) for i in h_boundary], []),
    )
end

function trotterized_layer_2d_v(Nx::Integer, Ny::Integer, istart::Integer; boundary::Bool=false)
    v = Iterators.flatten(repeat([i for i in 1:Nx], 1, Ny÷2-ifelse(iseven(Ny)&&iseven(istart), 1, 0)) .+ transpose([istart-1:2:Ny-2...].*Nx))
    v_boundary = [1:Nx...] .+ Nx*(Ny-1)
    L = Nx * Ny
    vcat(
        [circuit_gate(i, i+Nx, EntanglementXXGate(0.05π*randn())) for i in v],
        ifelse(boundary&&(isodd(Ny)==isodd(istart)), [circuit_gate(i, j, EntanglementXXGate(0.05π*randn())) for (i,j) in zip(v_boundary, 1:Nx)], []),
        [circuit_gate(i, i+Nx, EntanglementYYGate(0.05π*randn())) for i in v],
        ifelse(boundary&&(isodd(Ny)==isodd(istart)), [circuit_gate(i, j, EntanglementYYGate(0.05π*randn())) for (i,j) in zip(v_boundary, 1:Nx)], []),
        [circuit_gate(i, i+Nx, EntanglementZZGate(0.05π*randn())) for i in v],
        ifelse(boundary&&(isodd(Ny)==isodd(istart)), [circuit_gate(i, j, EntanglementZZGate(0.05π*randn())) for (i,j) in zip(v_boundary, 1:Nx)], [])
    )
end

function main()

    adj = lattice_adjacency_map(Nx, Ny; pbc=false)

    # number of neighbors
    sum(adj, dims=1)

    # number of lattice sites
    L = size(adj, 1)

    adj_horz = lattice_adjacency_map_horz(Nx, Ny; pbc=false)
    adj_vert = lattice_adjacency_map_vert(Nx, Ny; pbc=false)

    # display(adj_vert)
    # model parameters
    β = 2.6
    Jx = 1.0
    Jy = 0.6

    # construct Hamiltonian
    J = Jx * reshape(kron(adj_horz, [1, 1, 1]), (3, L, L)) + Jy * reshape(kron(adj_vert, [1, 1, 1]), (3, L, L))

    H = construct_hamiltonian(J)

    if debug
        @debug "Constructed Hamiltonian is equivalent to sample $(H ≈ H')"
        @assert(H ≈ H')
    end

    if visualize
        # show eigenvalues
        scatter(eigvals(Matrix(H)), xlabel=L"i", ylabel=L"\lambda_i", legend=false)
    end

    thermal_logZ(β*H)

    σtherm = thermal_state(β*H)
    # typeof(σtherm)

    # consistency check
    @assert tr(σtherm)  ≈ 1 "Trace of thermal state should be 1"

    if visualize
        heatmap(log.(max.(σtherm, 1e-4)), title=L"\sigma_{\beta}", yflip=true, aspect_ratio=:equal)
    end

    # parameterized quantum circuit gate chain
    cgc = CircuitGate[]
    layers = 8
    for _ in layers
        append!(cgc, single_qubit_layer(L))
        append!(cgc, trotterized_layer_1d(L, 1, boundary=isboundary))
        append!(cgc, trotterized_layer_1d(L, 2, boundary=isboundary))

        append!(cgc, single_qubit_layer(L))
        append!(cgc, trotterized_layer_2d_h(Nx, Ny, 1; boundary=isboundary))
        append!(cgc, trotterized_layer_2d_h(Nx, Ny, 2; boundary=isboundary))
        append!(cgc, single_qubit_layer(L))
        append!(cgc, trotterized_layer_2d_v(Nx, Ny, 1; boundary=isboundary))
        append!(cgc, trotterized_layer_2d_v(Nx, Ny, 2; boundary=isboundary))

    end
    append!(cgc, single_qubit_layer(L))

    # use Hamiltonian as measurement operator; note that H gets converted to a complex matrix here
    Hmeas = MeasurementOperator(H, Tuple(1:L));

    # parameterized quantum circuit
    circ = Circuit{L}(cgc, [Hmeas])

    # example
    latent_density([π/3, 3π/8])

    if visualize
        # visualize binary entropy
        plot(0:0.05:1, binary_entropy.(0:0.05:1), xlabel="p", ylabel="S(p)", legend=false)
    end

    # target function: β tr[ρ H] - S(ρ), with ρ = U ρlatent U†
    ftarget(θlist) = β * apply(latent_density(θlist), circ)[1] - sum(binary_entropy.((1 .+ cos.(θopt)) ./ 2))

    # initial random θ parameters, which are to be optimized
    θopt = π/2 .* (1 .+ 0.5*randn(L))
    # θopt = π/2 .* (6 .+ 3 .* randn(L))
    # # consistency check
    @assert neumann_entropy(matrix(apply(latent_density(θopt), cgc))) ≈ sum(binary_entropy.((1 .+ cos.(θopt)) ./ 2))
    @assert neumann_entropy(matrix(latent_density(θopt))) ≈ sum(binary_entropy.((1 .+ cos.(θopt)) ./ 2))
    f_neumann(θopt) = neumann_entropy(matrix(apply(latent_density(θopt), cgc)))
    f_binary(θopt) = sum(binary_entropy.((1 .+ cos.(θopt)) ./ 2))
    @assert all(isapprox.(ngradient(f_neumann, θopt), ngradient(f_binary, θopt), rtol=1e-5))

    #example
    @show ftarget(θopt)

    # gather parameters from circuit
    paras1 = Flux.params(circ)
    # paras2 = Flux.params(θopt)
    # measurement operator (Hamiltonian) not "trainable" here
    delete!(paras1, circ.meas[1].operator)
    # add θ parameters
    Flux.params!(paras1, θopt)

    # there is not actually any input data for training
    data = ncycle([()], 6000)

    # define optimizer
    opt = RMSProp()
    # opt = AdaMax()

    # define evaluation function
    evalcb() = println("ftarget(θopt): $(ftarget(θopt)), reference: $(thermal_logZ(β*H)); trdistance(θopt, σtherm, cgc): $(trdistance(θopt, σtherm, cgc))")

    # perform optimization

    Flux.train!(() -> ftarget(θopt), paras1, data, opt, cb=Flux.throttle(evalcb, 2))
    # println("alternating training")
    # Flux.train!(() -> ftarget(θopt), paras1, data, opt, cb=Flux.throttle(evalcb, 2))
    # println("alternating training")
    # Flux.train!(() -> ftarget(θopt), paras2, data, opt, cb=Flux.throttle(evalcb, 2))
    # println("alternating training")
    # Flux.train!(() -> ftarget(θopt), paras1, data, opt, cb=Flux.throttle(evalcb, 2))

    # seems like the Ansatz is too restricted, or the optimization trapped in a local minimum, to further descrease the distance
    @show trdistance(θopt, σtherm, cgc)

    # optimized variational density matrix
    ρopt = apply(latent_density(θopt), cgc)
    # println()
    # display(real.(matrix(ρopt)))
    # println()
    # display(σtherm)

    if visualize
        # visualize optimized density matrix
        heatmap(log.(max.(real(matrix(ρopt)), 1e-4)), title=L"\rho_{\mathrm{opt}}", yflip=true, aspect_ratio=:equal)
    end
end

main()

max_stars_repo_name: keefehuang/vqt

max_issues_repo_name: keefehuang/vqt

max_forks_repo_name: keefehuang/vqt


text:
# # Simple diffusion example
#
# This script provides our simplest example of Oceananigans.jl functionality:
# the diffusion of a one-dimensional Gaussian. This example demonstrates
#
#   * how to load `Oceananigans.jl`;
#   * how to instantiate an `Oceananigans.jl` `Model`;
#   * how to set an initial condition with a function;
#   * how to time-step a model forward, and finally
#   * how to look at results.
#
# ## Using `Oceananigans.jl`
#
# To use `Oceananigans.jl` after it has been installed, we bring
# `Oceananigans.jl` functions and names into our 'namespace' by writing

using Oceananigans

# We also use `PyPlot.jl` for plotting and `Printf` to format plot legends:

using PyPlot, Printf

# ## Instantiating and configuring a `Model`
#
# To begin using Oceananigans, we instantiate a `Model` by calling the
# `Model` constructor:

model = Model(
    grid = RegularCartesianGrid(size = (1, 1, 128), length = (1, 1, 1)),
    closure = ConstantIsotropicDiffusivity(κ = 1.0)
)

# The keyword arguments `grid` and `closure` indicate that
# our model grid is Cartesian with uniform grid spacing, that our diffusive
# stress and tracer fluxes are determined by diffusion with a constant
# diffusivity `κ` (note that we do not use viscosity in this example).

# Note that by default, a `Model` has no-flux boundary condition on all
# variables. Next, we set an initial condition on our "passive tracer",
# temperature. Our objective is to observe the diffusion of a Gaussian.

## Build a Gaussian initial condition function with width `δ`:
δ = 0.1
Tᵢ(x, y, z) = exp( -(z + 0.5)^2 / (2δ^2) )

## Set `model.tracers.T` to the function `Tᵢ`:
set!(model, T=Tᵢ)

# ## Running your first `Model`
#
# Finally, we time-step the model forward using the function
# `time_step!`, with a time-step size that ensures numerical stability.

## Time-scale for diffusion across a grid cell
cell_diffusion_time_scale = model.grid.Δz^2 / model.closure.κ.T

## The function `time_step!` executes `Nt` time steps with step size `Δt`
## using a second-order Adams-Bashforth method
time_step!(model, Nt = 1000, Δt = 0.1 * cell_diffusion_time_scale)

# ## Visualizing the results
#
# We use `PyPlot.jl` to look at the results.

## A convenient function for generating a label with the Current model time
tracer_label(model) = @sprintf("\$ t=%.3f \$", model.clock.time)

## Create a figure with `PyPlot.jl`
close("all")
fig, ax = subplots()
title("Diffusion of a Gaussian")
xlabel("Tracer concentration")
ylabel(L"z")

## Plot initial condition
plot(Tᵢ.(0, 0, model.grid.zC), model.grid.zC, "--", label=L"t=0")

## Plot current solution
plot(interior(model.tracers.T)[1, 1, :], model.grid.zC, label=tracer_label(model))
legend()
gcf()

# Interesting! Running the model even longer makes even more interesting results.

for i = 1:3
    time_step!(model, Nt = 1000, Δt = 0.1 * cell_diffusion_time_scale)
    plot(interior(model.tracers.T)[1, 1, :], model.grid.zC, label=tracer_label(model))
end

legend()
gcf()

ext: jl


text:
function StorelessCovariance(arg0::jint)
    return StorelessCovariance((jint,), arg0)
end

function StorelessCovariance(arg0::jint, arg1::jboolean)
    return StorelessCovariance((jint, jboolean), arg0, arg1)
end

function append(obj::StorelessCovariance, arg0::StorelessCovariance)
    return jcall(obj, "append", void, (StorelessCovariance,), arg0)
end

function get_covariance(obj::StorelessCovariance, arg0::jint, arg1::jint)
    return jcall(obj, "getCovariance", jdouble, (jint, jint), arg0, arg1)
end

function get_covariance_matrix(obj::StorelessCovariance)
    return jcall(obj, "getCovarianceMatrix", RealMatrix, ())
end

function get_data(obj::StorelessCovariance)
    return jcall(obj, "getData", Vector{Vector{jdouble}}, ())
end

function get_n(obj::StorelessCovariance)
    return jcall(obj, "getN", jint, ())
end

function increment(obj::StorelessCovariance, arg0::Vector{jdouble})
    return jcall(obj, "increment", void, (Vector{jdouble},), arg0)
end


text:
using OmegaGrav
using Dierckx
using Plots, LaTeXStrings
include("dodlnMh.jl")
# %% Compute the matter power spectrum using CLASS
include("compute_pk_class.jl")
# %% Compute dΩgrav/dlnMh and dΩtherm/dlnMh at z = 0
z = 0
# Define functions to return linear and non-linear power spectra
# Note: The CLASS code takes wavenumbers in units of 1/Mpc (no h) and
# return power spectra in units of Mpc^3 (no 1/h^3).
pkcb_class(kovh) = cosmo.pk_cb_lin(kovh * h0, z) * h0^3
# Spline interpolate in log(k)
lnk = log(1e-4):0.05:log(100)
pkcb = Spline1D(lnk, pkcb_class.(exp.(lnk)))
# %% Define a function to return dΩgrav/dlnMh from Halos, excluding the neutrino contribution
dΩgdlnMh(lnMh) = dogravdlnMh(x -> pkcb(log(x)), z, Ωm, lnMh, Ωcb)
# %% Define a function to return dΩtherm/dlnMh from Halos, excluding the neutrino contribution
dΩthdlnMh(lnMh) = dothermdlnMh(x -> pkcb(log(x)), z, Ωm, h0, lnMh, Ωcb)

#%% Plot results and save to figure2.pdf
fb = params["omega_b"] / (params["omega_cdm"] + params["omega_b"])
lnMh = log(1e11):0.1:log(5e15)
p = plot(
   exp.(lnMh),
   -2 * fb * dΩgdlnMh.(lnMh) / 3,
   xaxis = :log,
   yaxis = :log,
   xlab = "Halo Mass [M⊙/h]",
   ylab = L"d\Omega/d\ln M",
   lab = L"\Omega_W^{ref,halo}",
   legend = :topleft,
   legendfontsize = 12,
   labelfontsize = 15,
   lw = 5,
   c = :green,
   ls = :dot,
   ylims = [1e-10, 1e-8],
   xlims = [1e11, 5e15],
)
p = plot!(exp.(lnMh), dΩthdlnMh.(lnMh), lw = 5, c = 1, lab = L"\Omega_{th}")

# %% Compute dΩgrav/dlnMh and dΩtherm/dlnMh at z = 1
z = 1
pkcb_class(kovh) = cosmo.pk_cb_lin(kovh * h0, z) * h0^3
pkcb = Spline1D(lnk, pkcb_class.(exp.(lnk)))
dΩgdlnMh(lnMh) = dogravdlnMh(x -> pkcb(log(x)), z, Ωm, lnMh, Ωcb)
dΩthdlnMh(lnMh) = dothermdlnMh(x -> pkcb(log(x)), z, Ωm, h0, lnMh, Ωcb)
p = plot!(
   exp.(lnMh),
   -2 * fb * dΩgdlnMh.(lnMh) / 3,
   lw = 2,
   c = :green,
   ls = :dot,
   lab = "",
)
p = plot!(exp.(lnMh), dΩthdlnMh.(lnMh), lw = 2, c = 1, lab = "")
savefig("figure2.pdf")
display(p)

# %% Clean CLASS (the equivalent of the struct_free() in the `main`
# of CLASS. This step is primordial when running in a loop over different
# cosmologies, as you will saturate your memory very fast if you ommit
# it.
cosmo.struct_cleanup()
# If you want to change completely the cosmology, you should also
# clean the arguments, otherwise, if you are simply running on a loop
# of different values for the same parameters, this step is not needed
cosmo.empty()


text:
@testset "DQNLearner" begin
    env = CartPoleEnv(; T = Float32, seed = 11)
    ns, na = length(rand(get_observation_space(env))), length(get_action_space(env))
    agent = Agent(
        policy = QBasedPolicy(
            learner = DQNLearner(
                approximator = NeuralNetworkApproximator(
                    model = Chain(
                        Dense(ns, 128, relu; initW = seed_glorot_uniform(seed = 17)),
                        Dense(128, 128, relu; initW = seed_glorot_uniform(seed = 23)),
                        Dense(128, na; initW = seed_glorot_uniform(seed = 39)),
                    ) |> gpu,
                    optimizer = ADAM(),
                ),
                target_approximator = NeuralNetworkApproximator(
                    model = Chain(
                        Dense(ns, 128, relu; initW = seed_glorot_uniform(seed = 17)),
                        Dense(128, 128, relu; initW = seed_glorot_uniform(seed = 23)),
                        Dense(128, na; initW = seed_glorot_uniform(seed = 39)),
                    ) |> gpu,
                    optimizer = ADAM(),
                ),
                loss_func = huber_loss,
                stack_size = nothing,
                batch_size = 32,
                update_horizon = 1,
                min_replay_history = 100,
                update_freq = 1,
                target_update_freq = 100,
                seed = 22,
            ),
            explorer = EpsilonGreedyExplorer(
                kind = :exp,
                ϵ_stable = 0.01,
                decay_steps = 500,
                seed = 33,
            ),
        ),
        trajectory = CircularCompactSARTSATrajectory(
            capacity = 1000,
            state_type = Float32,
            state_size = (ns,),
        ),
    )
    hook = ComposedHook(TotalRewardPerEpisode(), TimePerStep())
    run(agent, env, StopAfterStep(10000), hook)

    @info "stats for DQNLearner" avg_reward = mean(hook[1].rewards) avg_fps =
        1 / mean(hook[2].times)
end


text:
# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    BlockPartition(sides)
    BlockPartition(side₁, side₂, ...)

A method for partitioning spatial objects into blocks of given `sides`.
"""
struct BlockPartition{Dim,T} <: PartitionMethod
  sides::SVector{Dim,T}
end

BlockPartition(sides::NTuple{Dim,T}) where {Dim,T} = BlockPartition{Dim,T}(sides)

BlockPartition(sides::Vararg{T,Dim}) where {Dim,T} = BlockPartition(sides)

function partition(object, method::BlockPartition)
  Dim = embeddim(object)
  T = coordtype(object)

  psides = method.sides
  bbox = boundingbox(object)

  @assert all(psides .≤ sides(bbox)) "invalid block sides"

  # bounding box properties
  lo, up = coordinates.(extrema(bbox))
  ce = coordinates(center(bbox))

  # find number of blocks to left and right
  nleft  = @. ceil(Int, (ce - lo) / psides)
  nright = @. ceil(Int, (up - ce) / psides)

  start   = @. ce - nleft * psides
  nblocks = @. nleft + nright

  subsets   = [Vector{Int}() for i in 1:prod(nblocks)]
  neighbors = [Vector{Int}() for i in 1:prod(nblocks)]

  # Cartesian to linear indices
  linear = LinearIndices(Dims(nblocks))

  coords = MVector{Dim,T}(undef)
  for j in 1:nelements(object)
    coordinates!(coords, object, j)

    # find block coordinates
    c = @. floor(Int, (coords - start) / psides) + 1
    @inbounds for i in 1:Dim
      c[i] = clamp(c[i], 1, nblocks[i])
    end
    bcoords = CartesianIndex(Tuple(c))

    # block index
    i = linear[bcoords]

    append!(subsets[i], j)
  end

  # neighboring blocks metadata
  bstart  = CartesianIndex(ntuple(i -> 1, Dim))
  boffset = CartesianIndex(ntuple(i -> 1, Dim))
  bfinish = CartesianIndex(Dims(nblocks))
  for (i, bcoords) in enumerate(bstart:bfinish)
    for b in (bcoords - boffset):(bcoords + boffset)
      if all(Tuple(bstart) .≤ Tuple(b) .≤ Tuple(bfinish)) && b ≠ bcoords
        push!(neighbors[i], linear[b])
      end
    end
  end

  # filter out empty blocks
  empty = isempty.(subsets)
  subsets = subsets[.!empty]
  neighbors = neighbors[.!empty]
  for i in findall(empty)
    for n in neighbors
      setdiff!(n, i)
    end
  end

  # save metadata
  metadata = Dict(:neighbors => neighbors)

  Partition(object, subsets, metadata)
end


text:
##### Arithmetics on contiguous ranks #####

for m=0:4, n=0:4
    global addrank
    @eval addrank(::Type{ContRank{$m}}, ::Type{ContRank{$n}}) = ContRank{$(m+n)}
end

addrank{M,N}(::Type{ContRank{M}}, ::Type{ContRank{N}}) = ContRank{M+N}
addrank{N}(::Type{ContRank{0}}, ::Type{ContRank{N}}) = ContRank{N}
addrank{N}(::Type{ContRank{N}}, ::Type{ContRank{0}}) = ContRank{N}

for m=0:4, n=0:4
    global minrank
    @eval minrank(::Type{ContRank{$m}}, ::Type{ContRank{$n}}) = ContRank{$(min(m,n))}
end

minrank{M,N}(::Type{ContRank{M}}, ::Type{ContRank{N}}) = ContRank{min(M,N)}
minrank{N}(::Type{ContRank{0}}, ::Type{ContRank{N}}) = ContRank{0}
minrank{N}(::Type{ContRank{N}}, ::Type{ContRank{0}}) = ContRank{0}

for m=0:4, n=0:4
    global restrict_crank
    @eval restrict_crank(::Type{ContRank{$m}}, ::NTuple{$n,Int}) = ContRank{$(min(m,n))}
end

restrict_crank{M,N}(::Type{ContRank{M}}, ::NTuple{N,Int}) = ContRank{min(M,N)}
restrict_crank{N}(::Type{ContRank{0}}, ::NTuple{N,Int}) = ContRank{0}
restrict_crank{N}(::Type{ContRank{N}}, ::NTuple{0,Int}) = ContRank{0}

### contiguous rank computation based on indices

# 0D
contrank() = ContRank{0}

# 1D
@compat contrank(i::Union{Colon,UnitRange}) = ContRank{1}
contrank(i::Subs) = ContRank{0}

# 2D
@compat contrank(i1::Colon, i2::Union{Colon,UnitRange}) = ContRank{2}
contrank(i1::Colon, i2::Subs) = ContRank{1}
contrank(i1::UnitRange, i2::Subs) = ContRank{1}
contrank(i1::Subs, i2::Subs) = ContRank{0}

# 3D
@compat contrank(i1::Colon, i2::Colon, i3::Union{Colon,UnitRange}) = ContRank{3}
contrank(i1::Colon, i2::Colon, i3::Subs) = ContRank{2}
contrank(i1::Colon, i2::UnitRange, i3::Subs) = ContRank{2}
contrank(i1::Colon, i2::Subs, i3::Subs) = ContRank{1}
contrank(i1::UnitRange, i2::Subs, i3::Subs) = ContRank{1}
contrank(i1::Subs, i2::Subs, i3::Subs) = ContRank{0}

# ND
contrank(i1::Colon, i2::Colon, i3::Colon, I::Subs...) = addrank(ContRank{3}, contrank(I...))
contrank(i1::Colon, i2::Colon, i3::UnitRange, I::Subs...) = ContRank{3}
contrank(i1::Colon, i2::Colon, i3::Subs, I::Subs...) = ContRank{2}
contrank(i1::Colon, i2::UnitRange, i3::Subs, I::Subs...) = ContRank{2}
contrank(i1::Colon, i2::Subs, i3::Subs, I::Subs...) = ContRank{1}
contrank(i1::UnitRange, i2::Subs, i3::Subs, I::Subs...) = ContRank{1}
contrank(i1::Subs, i2::Subs, i3::Subs, I::Subs...) = ContRank{0}

# contiguous rank with array & arrayviews

acontrank(a::Array, i1::Subs) = contrank(i1)
acontrank(a::Array, i1::Subs, i2::Subs) = contrank(i1, i2)
acontrank(a::Array, i1::Subs, i2::Subs, i3::Subs) = contrank(i1, i2, i3)
acontrank(a::Array, i1::Subs, i2::Subs, i3::Subs, i4::Subs) = contrank(i1, i2, i3, i4)
acontrank(a::Array, i1::Subs, i2::Subs, i3::Subs, i4::Subs, i5::Subs, I::Subs...) =
    contrank(i1, i2, i3, i4, i5, I...)

acontrank{T,N}(a::StridedArrayView{T,N,N}, i1::Subs) = contrank(i1)
acontrank{T,N}(a::StridedArrayView{T,N,N}, i1::Subs, i2::Subs) = contrank(i1, i2)
acontrank{T,N}(a::StridedArrayView{T,N,N}, i1::Subs, i2::Subs, i3::Subs) = contrank(i1, i2, i3)
acontrank{T,N}(a::StridedArrayView{T,N,N}, i1::Subs, i2::Subs, i3::Subs, i4::Subs) =
    contrank(i1, i2, i3, i4)
acontrank{T,N}(a::StridedArrayView{T,N,N}, i1::Subs, i2::Subs, i3::Subs, i4::Subs, i5::Subs, I::Subs...) =
    contrank(i1, i2, i3, i4, i5, I...)

acontrank{T,N,M}(a::StridedArrayView{T,N,M}, i1::Subs) = minrank(contrank(i1), ContRank{M})
acontrank{T,N,M}(a::StridedArrayView{T,N,M}, i1::Subs, i2::Subs) =
    minrank(contrank(i1, i2), ContRank{M})
acontrank{T,N,M}(a::StridedArrayView{T,N,M}, i1::Subs, i2::Subs, i3::Subs) =
    minrank(contrank(i1, i2, i3), ContRank{M})
acontrank{T,N,M}(a::StridedArrayView{T,N,M}, i1::Subs, i2::Subs, i3::Subs, i4::Subs) =
    minrank(contrank(i1, i2, i3, i4), ContRank{M})
acontrank{T,N,M}(a::StridedArrayView{T,N,M}, i1::Subs, i2::Subs, i3::Subs, i4::Subs, i5::Subs, I::Subs...) =
    minrank(contrank(i1, i2, i3, i4, i5, I...), ContRank{M})


text:
# Use baremodule to shave off a few KB from the serialized `.ji` file
baremodule Kissat_jll
using Base
using Base: UUID
import JLLWrappers

JLLWrappers.@generate_main_file_header("Kissat")
JLLWrappers.@generate_main_file("Kissat", UUID("971a2a11-349f-5eb6-a4bf-bc548fa92b3a"))
end  # module Kissat_jll


text:
"""
Author: Eric Ford
Created: September 2020
"""

#module FitTemplateToCCF

#import ..RVFromCCF: AbstractMeasureRvFromCCF, MeasureRvFromCCFQuadratic
# packages
using LinearAlgebra

# TODO: export once passes test that it's working
#export MeasureRvFromCCFTemplateNonDiagCovar

"""  Functor to estimate RV based on fitting a template and its derivative near minimum of the CCF. """
struct MeasureRvFromCCFTemplateNonDiagCovar <: AbstractMeasureRvFromCCF
    v_grid::AbstractVector{Float64}
    template::AbstractVector{Float64}
    deriv::AbstractVector{Float64}
	mean_var::Float64
	near_diag_covar::AbstractArray{Float64,2}
	v_idx_to_fit::UnitRange{Int64}
end

#=
# Todo: Replace with version in RvSpectMLBase or a (brute force?/temporal?) GP?
function numerical_deriv( x::AA1, y::AA2 )   where { T1<:Real, AA1<:AbstractArray{T1,1}, T2<:Real, AA2<:AbstractArray{T2,1}  }
	@assert length(x) == length(y)
	dydx = zeros(size(y))
	dydx[1] = (y[2]-y[1])/(x[2]-x[1])
	dydx[2:end-1] .= (y[3:end].-y[1:end-2])./(x[3:end].-x[1:end-2])
	dydx[end] = (y[end]-y[end-1])/(x[end]-x[end-1])
	return dydx
end
=#

"""
Construct functor to estimate RV based on the CCF.
TODO: Implement correctly.  Not yet working/tested.
Optional Arguments:
- `v_grid`:  list of velocities where template is evaluated
- `tempalte`: template ccf evaluated at v_grid
- `frac_of_width_to_fit`: (0.5)
- `measure_width_at_frac_depth`: (0.5)
"""
function MeasureRvFromCCFTemplateNonDiagCovar(; v_grid::AbstractVector{T1},
                                    template::AbstractVector{T2},
									frac_of_width_to_fit::Real = default_frac_of_width_to_fit,
                                    measure_width_at_frac_depth::Real = default_measure_width_at_frac_depth,
									mean_var::Real = 0,
									near_diag_covar::AbstractArray{T3,2} = zeros(length(v_grid),length(v_grid))
                                     ) where { T1<:Real, T2<:Real, T3<:Real }
    @assert length(v_grid) == length(template)
    @assert length(v_grid) >= 3
    @assert 0.25 <= measure_width_at_frac_depth <= 0.75
    @assert 0.1 <= frac_of_width_to_fit <= 5.0  # TODO: FIgure out appropriate range
	@assert 0 <= mean_var < Inf
	# Estimate derivative
    deriv = numerical_deriv(v_grid, template)
	# find the min and fit only the part near the minimum of the CCF
    v_min, v_idx_to_fit = find_idx_at_and_around_minimum(v_grid, template, frac_of_width_to_fit=frac_of_width_to_fit, measure_width_at_frac_depth=measure_width_at_frac_depth)
	#near_diag_covar = zeros(length(v_grid),length(v_grid))
    MeasureRvFromCCFTemplateNonDiagCovar(v_grid, template, deriv, mean_var, near_diag_covar, v_idx_to_fit)
end

function set_mean_var(in::MeasureRvFromCCFTemplateNonDiagCovar, mean_var::Real)
	MeasureRvFromCCFTemplateNonDiagCovar(v_grid=in.v_grid,template=in.template, frac_of_width_to_fit=in.frac_of_width_to_fit, measure_width_at_frac_depth=in.measure_width_at_frac_depth,
		mean_var=mean_var, near_diag_covar=in.near_diag_covar )
end

function (mrv::MeasureRvFromCCFTemplateNonDiagCovar)(vels::A1, ccf::A2, ccf_var::A3, ccf_covar::A4  ) where {T1<:Real, A1<:AbstractArray{T1,1}, T2<:Real, A2<:AbstractArray{T2,1}, T3<:Real, A3<:AbstractArray{T3,1}, T4<:Real, A4<:AbstractArray{T4,2}  }
		@assert all(vels .== mrv.v_grid)  # Could relax this, but keep it simple for now
		# fit only the part near the minimum of the CCF
		deriv = view(mrv.deriv,mrv.v_idx_to_fit)
		diag_scale_factor = mean(ccf_var)-mrv.mean_var
		covar = view(ccf_covar,mrv.v_idx_to_fit,mrv.v_idx_to_fit) + diag_scale_factor*view(mrv.near_diag_covar,mrv.v_idx_to_fit,mrv.v_idx_to_fit)
        # find the min and fit only the part near the minimum of the CCF
		ccf_inv_covar_times_deriv = covar \ deriv
		denom = deriv' * ccf_inv_covar_times_deriv
		# Minus sign due to weird convention of defining CCF to look like an absorption line (e.g., 1 at Inf minimum at best fit velocity)
		rv = -(view(ccf,mrv.v_idx_to_fit) .- view(mrv.template,mrv.v_idx_to_fit) )' * ccf_inv_covar_times_deriv
		rv *= (1/denom)
		σ_rv = sqrt(1/denom)
		return (rv=rv, σ_rv=σ_rv)
end

max_stars_repo_name: RvSpectML/EchelleCCFs

max_issues_repo_name: RvSpectML/EchelleCCFs

max_forks_repo_name: RvSpectML/EchelleCCFs


text:
using ClutteredEnvPathOpt
using LightGraphs
import Polyhedra
import GLPK
using Pipe
using Plots
using JuMP, Gurobi, PiecewiseLinearOpt

function get_M_A_b_easy(obstacles)
    Ms = []
    As = []
    bs = []
    acc = [0] # accumulator # try using reduce() instead

    _, points, _, _, free_faces = ClutteredEnvPathOpt.construct_graph(obstacles)

    for face in free_faces
        v = Polyhedra.convexhull(map(i -> collect(points[i]), face)...)
        polygon = Polyhedra.polyhedron(
            v,
            Polyhedra.DefaultLibrary{Rational{Int64}}(Gurobi.Optimizer)
        )
        halfspaces = Polyhedra.hrep(polygon).halfspaces
        temp = acc[end]
        push!(acc, temp + length(halfspaces))

        # for i in v.points
        #     println("$i $face")
        # end

        for hs in halfspaces
            A = hs.a'
            b = hs.β

            M = maximum(map(x -> A * x, v.points))

            push!(Ms, M)
            push!(As, A)
            push!(bs, b)
        end
    end
    # acc = acc[2:end]

    return vcat(Ms...), vcat(As...), vcat(bs...), acc
end

function get_M_A_b(points, free_faces)
    Ms = []
    As = []
    bs = []
    acc = [0] # try using reduce() instead

    # _, points, _, _, free_faces = ClutteredEnvPathOpt.construct_graph(obstacles)

    # Solve over the unitcell
    u = Polyhedra.convexhull([0//1,0//1],[0//1,1//1],[1//1,0//1],[1//1,1//1])
    unitcell = Polyhedra.polyhedron(u)

    for face in free_faces
        v = Polyhedra.convexhull(map(i -> collect(points[i]), face)...)
        polygon = Polyhedra.polyhedron(
            v,
            Polyhedra.DefaultLibrary{Rational{Int64}}(Gurobi.Optimizer)
        )
        halfspaces = Polyhedra.hrep(polygon).halfspaces
        temp = acc[end]
        push!(acc, temp + length(halfspaces))

        for hs in halfspaces
            A = hs.a'
            b = hs.β

            sub_model = JuMP.Model(Gurobi.Optimizer)
            #sub_x = JuMP.@variable(sub_model, [1:2])
            JuMP.@variable(sub_model, sub_x[1:2])
            JuMP.@constraint(sub_model, sub_x in unitcell)
            #JuMP.@constraint(sub_model, sub_x in safe_regions)
            JuMP.@objective(sub_model, Max, A * sub_x)

            JuMP.optimize!(sub_model)

            #return value.(sub_x)
            M = JuMP.objective_value(sub_model)

            push!(Ms, M)
            push!(As, A)
            push!(bs, b)
        end

        # A = vcat([hs.a for hs in halfspaces]...)
        # b = vcat([hs.β for hs in halfspaces]...)

        # M = map(hs ->
        #     begin
        #         A = hs.a
        #         b = hs.β

        #         sub_model = JuMP.Model(Gurobi.Optimizer)
        #         #sub_x = JuMP.@variable(sub_model, [1:2])
        #         JuMP.@variable(sub_model, sub_x[1:2])
        #         JuMP.@constraint(sub_model, sub_x in polygon)   # TODO "no method matching copy" error WHY DOES THIS ONLY WORK THE FIRST TIME
        #         JuMP.@objective(sub_model, Max, A' * sub_x)

        #         stat = JuMP.optimize!(sub_model)

        #         #return value.(sub_x)
        #         return JuMP.objective_value(sub_model)
        #     end,
        #     halfspaces
        # )

        # push!(Ms, M)
        # push!(As, A)
        # push!(bs, b)
    end
    # acc = acc[2:end]

    return vcat(Ms...), vcat(As...), vcat(bs...), acc
    #return Ms, As, bs, acc
end

# obstacles <- obstacles (list of polyhedra)
# N <- max number of steps (scalar)
# f1 <- initial left foot ([x, y, theta])
# f2 <- initial right foot ([x, y, theta])
# g <- goal pose
# Q_g <- pose weights to goal (4x4 mat)
# Q_r <- pose weights between steps (4x4 mat)
# q_t <- weight on unused steps (scalar)
# method <- "merged" for the compact biclique cover, "full" for the original biclique cover, "bigM" for big-M constraints
# d1 = 0.2 <- radius of reference foot circle
# d2 = 0.2 <- radius of moving foot circle
# p1 = [0, 0.07] <- center of reference foot circle
# p2 = [0, -0.27] <- center of moving foot circle
# delta_x_y_max = 0.1 <- max stride norm in space
# delta_θ_max = pi/4 <- max difference in θ
# L = 5 <- number of pieces of pwl sin/cos (scalar)
function solve_deits(obstacles, N, f1, f2, g, Q_g, Q_r, q_t; method="merged", d1=0.2, d2=0.2, p1=[0, 0.07], p2=[0, -0.27])

    _, points, graph, _, free_faces = ClutteredEnvPathOpt.construct_graph(obstacles)
    skeleton = LabeledGraph(graph)

    # model = JuMP.Model(JuMP.optimizer_with_attributes(Gurobi.Optimizer))
    model = JuMP.Model(JuMP.optimizer_with_attributes(Gurobi.Optimizer, "MIPGap" => .01, "TimeLimit" => 180))
    # model = JuMP.Model(JuMP.optimizer_with_attributes(Gurobi.Optimizer, "Heuristics"=> 0, "Cuts"=> 0, "Precrush"=>1, "MIPGap" => .01, "TimeLimit" => 180))

    # model has scalar variables x, y, θ, binary variable t
    JuMP.@variable(model, x[1:N])
    JuMP.@variable(model, y[1:N])
    JuMP.@variable(model, θ[1:N])
    JuMP.@variable(model, t[1:N], Bin)

    # Objective
    f = vec([[x[j], y[j], θ[j]] for j in 1:N])
    # f_no_theta = vec([[x[j], y[j]] for j in 1:N])
    JuMP.@objective(
        model,
        Min,
        ((f[N] - g)' * Q_g * (f[N] - g)) + sum(q_t * t) + sum((f[j + 1] - f[j])' * Q_r * (f[j + 1] - f[j]) for j in 1:(N-1))
    )

    cover = Set{Pair{Set{Int64}, Set{Int64}}}()
    merged_cover = Set{Pair{Set{Int64}, Set{Int64}}}()
    num_free_face_ineq = 0

    if method != "bigM"
        cover = ClutteredEnvPathOpt.find_biclique_cover(skeleton, free_faces)
        feg = ClutteredEnvPathOpt._find_finite_element_graph(skeleton, ClutteredEnvPathOpt._find_face_pairs(free_faces))
        merged_cover = ClutteredEnvPathOpt.biclique_merger(cover, feg)

        (valid_cover, _, _, _) = ClutteredEnvPathOpt._is_valid_biclique_cover_diff(feg, cover)
        (valid_cover_merged, _, _, _) = ClutteredEnvPathOpt._is_valid_biclique_cover_diff(feg, merged_cover)
    end

    # Footstep location constraints
    if method == "merged" && valid_cover_merged
        J = LightGraphs.nv(skeleton.graph)
        JuMP.@variable(model, λ[1:N, 1:J] >= 0)
        JuMP.@variable(model, z[1:N, 1:length(merged_cover)], Bin)
        for i in 1:N
            for (j,(A,B)) in enumerate(merged_cover)
                JuMP.@constraint(model, sum(λ[i, v] for v in A) <= z[i, j])
                JuMP.@constraint(model, sum(λ[i, v] for v in B) <= 1 - z[i, j])
            end
            JuMP.@constraint(model, sum(λ[i,:]) == 1)
            JuMP.@constraint(model, x[i] == sum(points[j].first * λ[i, j] for j in 1:J))
            JuMP.@constraint(model, y[i] == sum(points[j].second * λ[i, j] for j in 1:J))
        end
    elseif method != "bigM" && valid_cover
        J = LightGraphs.nv(skeleton.graph)
        JuMP.@variable(model, λ[1:N, 1:J] >= 0)
        JuMP.@variable(model, z[1:N, 1:length(cover)], Bin)
        for i in 1:N
            for (j,(A,B)) in enumerate(cover)
                JuMP.@constraint(model, sum(λ[i, v] for v in A) <= z[i, j])
                JuMP.@constraint(model, sum(λ[i, v] for v in B) <= 1 - z[i, j])
            end
            JuMP.@constraint(model, sum(λ[i,:]) == 1)
            JuMP.@constraint(model, x[i] == sum(points[j].first * λ[i, j] for j in 1:J))
            JuMP.@constraint(model, y[i] == sum(points[j].second * λ[i, j] for j in 1:J))
        end
    else
        # Runs if method = "bigM", or if a valid biclique cover is not found
        M, A, b, acc = get_M_A_b(points, free_faces)
        num_free_face_ineq = length(M)
        JuMP.@variable(model, z[1:N, 1:length(free_faces)], Bin)
        for j in 1:N
            for r in 1:length(free_faces)
                ids = (acc[r]+1):(acc[r+1])
                for i in ids
                    # JuMP.@constraint(model, A[r, :]' * [x[j], y[j]] <= b[r] * z[j, r] + M[r] * (1 - z[j, r]))
                    # JuMP.@constraint(model, A[r, :]' * [x[j], y[j]] <= b[r] * z[j, r] + 1 * (1 - z[j, r]))
                    # println("$i")
                    JuMP.@constraint(model, A[i, :]' * [x[j]; y[j]] <= b[i] * z[j, r] + M[i] * (1 - z[j, r]))
                end
            end
            JuMP.@constraint(model, sum(z[j, :]) == 1)
        end
    end

    # Reachability
    # s = [piecewiselinear(model, θ[j], range(0, stop=(2 * pi), length=L), sin) for j in 1:N]
    # c = [piecewiselinear(model, θ[j], range(0, stop=(2 * pi), length=L), cos) for j in 1:N]
    s_break_pts = [0, 5pi/16, 11pi/16, 21pi/16, 27pi/16, 2pi]
    c_break_pts = [0, 3pi/16, 13pi/16, 19pi/16, 29pi/16, 2pi]
    s = [piecewiselinear(model, θ[j], s_break_pts, sin) for j in 1:N]
    c = [piecewiselinear(model, θ[j], c_break_pts, cos) for j in 1:N]

    # For footstep j, the cirles are created from the frame of reference of footstep j-1
    # Use negative if in frame of reference of right foot (j = 1 is left foot)
    # Position is with respect to θ = 0
    for j in 2:N
        if j % 2 == 1
            # In frame of reference of right foot
            JuMP.@constraint(
                model,
                [
                    d1,
                    x[j] - x[j - 1] - c[j-1] * (-p1[1]) + s[j-1] * (-p1[2]),
                    y[j] - y[j - 1] - s[j-1] * (-p1[1]) - c[j-1] * (-p1[2])
                ] in SecondOrderCone()
            )

            JuMP.@constraint(
                model,
                [
                    d2,
                    x[j] - x[j - 1] - c[j-1] * (-p2[1]) + s[j-1] * (-p2[2]),
                    y[j] - y[j - 1] - s[j-1] * (-p2[1]) - c[j-1] * (-p2[2])
                ] in SecondOrderCone()
            )
        else
            # In frame of reference left foot
            JuMP.@constraint(
                model,
                [
                    d1,
                    x[j] - x[j - 1] - c[j-1] * p1[1] + s[j-1] * p1[2],
                    y[j] - y[j - 1] - s[j-1] * p1[1] - c[j-1] * p1[2]
                ] in SecondOrderCone()
            )

            JuMP.@constraint(
                model,
                [
                    d2,
                    x[j] - x[j - 1] - c[j-1] * p2[1] + s[j-1] * p2[2],
                    y[j] - y[j - 1] - s[j-1] * p2[1] - c[j-1] * p2[2]
                ] in SecondOrderCone()
            )
        end
    end

    # Set T to punish extra steps
    for j in 1:N
        if j % 2 == 1
            JuMP.@constraint(model, t[j] => {x[j] == f1[1]}) # use f1 and f2?
            JuMP.@constraint(model, t[j] => {y[j] == f1[2]})
            JuMP.@constraint(model, t[j] => {θ[j] == f1[3]})
        else
            JuMP.@constraint(model, t[j] => {x[j] == f2[1]})
            JuMP.@constraint(model, t[j] => {y[j] == f2[2]})
            JuMP.@constraint(model, t[j] => {θ[j] == f2[3]})
        end
    end

    # Max step distance
    # for j in 3:2:(N-1)
    #     JuMP.@constraint(model, [delta_x_y_max; f_no_theta[j] - f_no_theta[j - 2]] in SecondOrderCone())
    #     JuMP.@constraint(model, [delta_x_y_max; f_no_theta[j + 1] - f_no_theta[j - 1]] in SecondOrderCone())
    # end

    # Max theta difference for individual feet
    # for j in 3:2:(N-1)
    #     JuMP.@constraint(model, -delta_θ_max <= θ[j] - θ[j - 2] <= delta_θ_max)
    #     JuMP.@constraint(model, -delta_θ_max <= θ[j + 1] - θ[j - 1] <= delta_θ_max)
    # end

    # Max θ difference between feet
    for j in 2:N
        JuMP.@constraint(model, -pi/8 <= θ[j] - θ[j-1] <= pi/8)
    end

    # Initial footstep positions
    JuMP.@constraint(model, f[1] .== f1)
    JuMP.@constraint(model, f[2] .== f2)

    # Solve
    JuMP.optimize!(model)

    if method == "merged" && valid_cover_merged
        println("\n\nUsed merged cover.\n\n")
    elseif method != "bigM" && valid_cover
        println("\n\nUsed full cover.\n\n")
        method = "full"
    else
        println("\n\nUsed big-M constraints.\n\n")
        method = "bigM"
    end

    stats = (termination_status(model), solve_time(model), relative_gap(model), simplex_iterations(model),
            node_count(model), LightGraphs.nv(skeleton.graph), length(merged_cover), length(cover), length(free_faces), num_free_face_ineq, method)

    return value.(x), value.(y), value.(θ), value.(t), stats
end

function plot_steps(obstacles, x, y, θ)
    plot()
    ClutteredEnvPathOpt.plot_field(obstacles);
    scatter!(x[1:2:end], y[1:2:end], color="red", series_annotations=([Plots.text(string(x), :right, 8, "courier") for x in 1:2:length(x)]));
    scatter!(x[2:2:end], y[2:2:end], color="blue", series_annotations=([Plots.text(string(x), :right, 8, "courier") for x in 2:2:length(x)]));
    quiver!(x, y, quiver=(0.075 * cos.(θ), 0.075 * sin.(θ)))
    display(plot!(title="Footsteps"))
end

function plot_circles(x, y, theta; R1=0.20, R2=0.20, p1=[0, 0.07], p2=[0, -0.27])
    angles = LinRange(0,2*pi, 100)
    #plot()
    for i = 2:length(x)
        plot()
        if i % 2 == 1
            circlex1 = R1*cos.(angles) .+ ( x[i-1] + cos(theta[i-1])*(-p1[1]) - sin(theta[i-1])*(-p1[2]) )
            circley1 = R1*sin.(angles) .+ ( y[i-1] + sin(theta[i-1])*(-p1[1]) + cos(theta[i-1])*(-p1[2]) )

            circlex2 = R2*cos.(angles) .+ ( x[i-1] + cos(theta[i-1])*(-p2[1]) - sin(theta[i-1])*(-p2[2]) )
            circley2 = R2*sin.(angles) .+ ( y[i-1] + sin(theta[i-1])*(-p2[1]) + cos(theta[i-1])*(-p2[2]) )
            #scatter!([x[i-1]; x[i]], [y[i-1]; y[i]])
            scatter!([x[i-1]], [y[i-1]], color = "blue")
            scatter!([x[i]], [y[i]], color = "red")
            quiver!([x[i-1]; x[i]], [y[i-1]; y[i]],
                quiver=([0.075 * cos(theta[i-1]); 0.075 * cos(theta[i])],
                [0.075 * sin(theta[i-1]); 0.075 * sin(theta[i])])
                )
        else
            circlex1 = R1*cos.(angles) .+ ( x[i-1] + cos(theta[i-1])*(p1[1]) - sin(theta[i-1])*(p1[2]) )
            circley1 = R1*sin.(angles) .+ ( y[i-1] + sin(theta[i-1])*(p1[1]) + cos(theta[i-1])*(p1[2]) )

            circlex2 = R2*cos.(angles) .+ ( x[i-1] + cos(theta[i-1])*(p2[1]) - sin(theta[i-1])*(p2[2]) )
            circley2 = R2*sin.(angles) .+ ( y[i-1] + sin(theta[i-1])*(p2[1]) + cos(theta[i-1])*(p2[2]) )
            #scatter!([x[i-1]; x[i]], [y[i-1]; y[i]])
            scatter!([x[i-1]], [y[i-1]], color = "red")
            scatter!([x[i]], [y[i]], color = "blue")
            quiver!([x[i-1]; x[i]], [y[i-1]; y[i]],
                quiver=([0.075 * cos(theta[i-1]); 0.075 * cos(theta[i])],
                [0.075 * sin(theta[i-1]); 0.075 * sin(theta[i])])
                )
        end
        plot!(circlex1, circley1, color="dodgerblue")
        plot!(circlex2, circley2, color="maroon")
        display(plot!(legend=false, xlims=(-0.1,1.1), ylims=(-0.1,1.1)))
    end
    # display(plot!(legend=false, xlims=(-0.1,1.1), ylims=(-0.1,1.1)))
end


text:
"""
    OGR_SM_Create(OGRStyleTableH hStyleTable) -> OGRStyleMgrH
OGRStyleMgr factory.
### Parameters
* **hStyleTable**: pointer to OGRStyleTable or NULL if not working with a style table.
### Returns
an handle to the new style manager object.
"""
function create(hStyleTable::Ptr{OGRStyleTableH})
    checknull(ccall((:OGR_SM_Create,libgdal),Ptr{OGRStyleMgrH},(Ptr{OGRStyleTableH},),hStyleTable))
end

"""
    OGR_SM_Destroy(OGRStyleMgrH hSM) -> void
Destroy Style Manager.
### Parameters
* **hSM**: handle to the style manager to destroy.
"""
function destroy(hSM::Ptr{OGRStyleMgrH})
    ccall((:OGR_SM_Destroy,libgdal),Void,(Ptr{OGRStyleMgrH},),hSM)
end

"""
    OGR_SM_InitFromFeature(OGRStyleMgrH hSM,
                           OGRFeatureH hFeat) -> const char *
Initialize style manager from the style string of a feature.
### Parameters
* **hSM**: handle to the style manager.
* **hFeat**: handle to the new feature from which to read the style.
### Returns
a reference to the style string read from the feature, or NULL in case of error.
"""
function initfromfeature(hSM::Ptr{OGRStyleMgrH},hFeat::Ptr{OGRFeatureH})
    bytestring(ccall((:OGR_SM_InitFromFeature,libgdal),Cstring,(Ptr{OGRStyleMgrH},Ptr{OGRFeatureH}),hSM,hFeat))
end

"""
    OGR_SM_InitStyleString(OGRStyleMgrH hSM,
                           const char * pszStyleString) -> int
Initialize style manager from the style string.
### Parameters
* **hSM**: handle to the style manager.
* **pszStyleString**: the style string to use (can be NULL).
### Returns
TRUE on success, FALSE on errors.
"""
function initstylestring(hSM::Ptr{OGRStyleMgrH},pszStyleString)
    ccall((:OGR_SM_InitStyleString,libgdal),Cint,(Ptr{OGRStyleMgrH},Cstring),hSM,pszStyleString)
end

"""
    OGR_SM_GetPartCount(OGRStyleMgrH hSM,
                        const char * pszStyleString) -> int
Get the number of parts in a style.
### Parameters
* **hSM**: handle to the style manager.
* **pszStyleString**: (optional) the style string on which to operate. If NULL then the current style string stored in the style manager is used.
### Returns
the number of parts (style tools) in the style.
"""
function getpartcount(hSM::Ptr{OGRStyleMgrH},pszStyleString)
    ccall((:OGR_SM_GetPartCount,libgdal),Cint,(Ptr{OGRStyleMgrH},Cstring),hSM,pszStyleString)
end

"""
    OGR_SM_GetPart(OGRStyleMgrH hSM,
                   int nPartId,
                   const char * pszStyleString) -> OGRStyleToolH
Fetch a part (style tool) from the current style.
### Parameters
* **hSM**: handle to the style manager.
* **nPartId**: the part number (0-based index).
* **pszStyleString**: (optional) the style string on which to operate. If NULL then the current style string stored in the style manager is used.
### Returns
OGRStyleToolH of the requested part (style tools) or NULL on error.
"""
function getpart(hSM::Ptr{OGRStyleMgrH},nPartId::Integer,pszStyleString)
    checknull(ccall((:OGR_SM_GetPart,libgdal),Ptr{OGRStyleToolH},(Ptr{OGRStyleMgrH},Cint,Cstring),hSM,nPartId,pszStyleString))
end

"""
    OGR_SM_AddPart(OGRStyleMgrH hSM,
                   OGRStyleToolH hST) -> int
Add a part (style tool) to the current style.
### Parameters
* **hSM**: handle to the style manager.
* **hST**: the style tool defining the part to add.
### Returns
TRUE on success, FALSE on errors.
"""
function addpart(hSM::Ptr{OGRStyleMgrH},hST::Ptr{OGRStyleToolH})
    ccall((:OGR_SM_AddPart,libgdal),Cint,(Ptr{OGRStyleMgrH},Ptr{OGRStyleToolH}),hSM,hST)
end

"""
    OGR_SM_AddStyle(OGRStyleMgrH hSM,
                    const char * pszStyleName,
                    const char * pszStyleString) -> int
Add a style to the current style table.
### Parameters
* **hSM**: handle to the style manager.
* **pszStyleName**: the name of the style to add.
* **pszStyleString**: the style string to use, or NULL to use the style stored in the manager.
### Returns
TRUE on success, FALSE on errors.
"""
function addstyle(hSM::Ptr{OGRStyleMgrH},pszStyleName,pszStyleString)
    ccall((:OGR_SM_AddStyle,libgdal),Cint,(Ptr{OGRStyleMgrH},Cstring,Cstring),hSM,pszStyleName,pszStyleString)
end

"""
    OGR_ST_Create(OGRSTClassId eClassId) -> OGRStyleToolH
OGRStyleTool factory.
### Parameters
* **eClassId**: subclass of style tool to create. One of OGRSTCPen (1), OGRSTCBrush (2), OGRSTCSymbol (3) or OGRSTCLabel (4).
### Returns
an handle to the new style tool object or NULL if the creation failed.
"""
function create(eClassId::OGRSTClassId)
    checknull(ccall((:OGR_ST_Create,libgdal),Ptr{OGRStyleToolH},(OGRSTClassId,),eClassId))
end

"""
    OGR_ST_Destroy(OGRStyleToolH hST) -> void
Destroy Style Tool.
### Parameters
* **hST**: handle to the style tool to destroy.
"""
function destroy(hST::Ptr{OGRStyleToolH})
    ccall((:OGR_ST_Destroy,libgdal),Void,(Ptr{OGRStyleToolH},),hST)
end

"""
    OGR_ST_GetType(OGRStyleToolH hST) -> OGRSTClassId
Determine type of Style Tool.
### Parameters
* **hST**: handle to the style tool.
### Returns
the style tool type, one of OGRSTCPen (1), OGRSTCBrush (2), OGRSTCSymbol (3) or OGRSTCLabel (4). Returns OGRSTCNone (0) if the OGRStyleToolH is invalid.
"""
function gettype(hST::Ptr{OGRStyleToolH})
    ccall((:OGR_ST_GetType,libgdal),OGRSTClassId,(Ptr{OGRStyleToolH},),hST)
end

"""
    OGR_ST_GetUnit(OGRStyleToolH hST) -> OGRSTUnitId
Get Style Tool units.
### Parameters
* **hST**: handle to the style tool.
### Returns
the style tool units.
"""
function getunit(hST::Ptr{OGRStyleToolH})
    ccall((:OGR_ST_GetUnit,libgdal),OGRSTUnitId,(Ptr{OGRStyleToolH},),hST)
end

"""
    OGR_ST_SetUnit(OGRStyleToolH hST,
                   OGRSTUnitId eUnit,
                   double dfGroundPaperScale) -> void
Set Style Tool units.
### Parameters
* **hST**: handle to the style tool.
* **eUnit**: the new unit.
* **dfGroundPaperScale**: ground to paper scale factor.
"""
function setunit(hST::Ptr{OGRStyleToolH},eUnit::OGRSTUnitId,dfGroundPaperScale::Real)
    ccall((:OGR_ST_SetUnit,libgdal),Void,(Ptr{OGRStyleToolH},OGRSTUnitId,Cdouble),hST,eUnit,dfGroundPaperScale)
end

"""
    OGR_ST_GetParamStr(OGRStyleToolH hST,
                       int eParam,
                       int * bValueIsNull) -> const char *
Get Style Tool parameter value as string.
### Parameters
* **hST**: handle to the style tool.
* **eParam**: the parameter id from the enumeration corresponding to the type of this style tool (one of the OGRSTPenParam, OGRSTBrushParam, OGRSTSymbolParam or OGRSTLabelParam enumerations)
* **bValueIsNull**: pointer to an integer that will be set to TRUE or FALSE to indicate whether the parameter value is NULL.
### Returns
the parameter value as string and sets bValueIsNull.
"""
function getparamstr(hST::Ptr{OGRStyleToolH},eParam::Integer,bValueIsNull)
    bytestring(ccall((:OGR_ST_GetParamStr,libgdal),Cstring,(Ptr{OGRStyleToolH},Cint,Ptr{Cint}),hST,eParam,bValueIsNull))
end

"""
    OGR_ST_GetParamNum(OGRStyleToolH hST,
                       int eParam,
                       int * bValueIsNull) -> int
Get Style Tool parameter value as an integer.
### Parameters
* **hST**: handle to the style tool.
* **eParam**: the parameter id from the enumeration corresponding to the type of this style tool (one of the OGRSTPenParam, OGRSTBrushParam, OGRSTSymbolParam or OGRSTLabelParam enumerations)
* **bValueIsNull**: pointer to an integer that will be set to TRUE or FALSE to indicate whether the parameter value is NULL.
### Returns
the parameter value as integer and sets bValueIsNull.
"""
function getparamnum(hST::Ptr{OGRStyleToolH},eParam::Integer,bValueIsNull)
    ccall((:OGR_ST_GetParamNum,libgdal),Cint,(Ptr{OGRStyleToolH},Cint,Ptr{Cint}),hST,eParam,bValueIsNull)
end

"""
    OGR_ST_GetParamDbl(OGRStyleToolH hST,
                       int eParam,
                       int * bValueIsNull) -> double
Get Style Tool parameter value as a double.
### Parameters
* **hST**: handle to the style tool.
* **eParam**: the parameter id from the enumeration corresponding to the type of this style tool (one of the OGRSTPenParam, OGRSTBrushParam, OGRSTSymbolParam or OGRSTLabelParam enumerations)
* **bValueIsNull**: pointer to an integer that will be set to TRUE or FALSE to indicate whether the parameter value is NULL.
### Returns
the parameter value as double and sets bValueIsNull.
"""
function getparamdbl(hST::Ptr{OGRStyleToolH},eParam::Integer,bValueIsNull)
    ccall((:OGR_ST_GetParamDbl,libgdal),Cdouble,(Ptr{OGRStyleToolH},Cint,Ptr{Cint}),hST,eParam,bValueIsNull)
end

"""
    OGR_ST_SetParamStr(OGRStyleToolH hST,
                       int eParam,
                       const char * pszValue) -> void
Set Style Tool parameter value from a string.
### Parameters
* **hST**: handle to the style tool.
* **eParam**: the parameter id from the enumeration corresponding to the type of this style tool (one of the OGRSTPenParam, OGRSTBrushParam, OGRSTSymbolParam or OGRSTLabelParam enumerations)
* **pszValue**: the new parameter value
"""
function setparamstr(hST::Ptr{OGRStyleToolH},eParam::Integer,pszValue)
    ccall((:OGR_ST_SetParamStr,libgdal),Void,(Ptr{OGRStyleToolH},Cint,Cstring),hST,eParam,pszValue)
end

"""
    OGR_ST_SetParamNum(OGRStyleToolH hST,
                       int eParam,
                       int nValue) -> void
Set Style Tool parameter value from an integer.
### Parameters
* **hST**: handle to the style tool.
* **eParam**: the parameter id from the enumeration corresponding to the type of this style tool (one of the OGRSTPenParam, OGRSTBrushParam, OGRSTSymbolParam or OGRSTLabelParam enumerations)
* **nValue**: the new parameter value
"""
function setparamnum(hST::Ptr{OGRStyleToolH},eParam::Integer,nValue::Integer)
    ccall((:OGR_ST_SetParamNum,libgdal),Void,(Ptr{OGRStyleToolH},Cint,Cint),hST,eParam,nValue)
end

"""
    OGR_ST_SetParamDbl(OGRStyleToolH hST,
                       int eParam,
                       double dfValue) -> void
Set Style Tool parameter value from a double.
### Parameters
* **hST**: handle to the style tool.
* **eParam**: the parameter id from the enumeration corresponding to the type of this style tool (one of the OGRSTPenParam, OGRSTBrushParam, OGRSTSymbolParam or OGRSTLabelParam enumerations)
* **dfValue**: the new parameter value
"""
function setparamdbl(hST::Ptr{OGRStyleToolH},eParam::Integer,dfValue::Real)
    ccall((:OGR_ST_SetParamDbl,libgdal),Void,(Ptr{OGRStyleToolH},Cint,Cdouble),hST,eParam,dfValue)
end

"""
    OGR_ST_GetStyleString(OGRStyleToolH hST) -> const char *
Get the style string for this Style Tool.
### Parameters
* **hST**: handle to the style tool.
### Returns
the style string for this style tool or "" if the hST is invalid.
"""
function getstylestring(hST::Ptr{OGRStyleToolH})
    bytestring(ccall((:OGR_ST_GetStyleString,libgdal),Cstring,(Ptr{OGRStyleToolH},),hST))
end

"""
    OGR_ST_GetRGBFromString(OGRStyleToolH hST,
                            const char * pszColor,
                            int * pnRed,
                            int * pnGreen,
                            int * pnBlue,
                            int * pnAlpha) -> int
Return the r,g,b,a components of a color encoded in #RRGGBB[AA] format.
### Parameters
* **hST**: handle to the style tool.
* **pszColor**: the color to parse
* **pnRed**: pointer to an int in which the red value will be returned
* **pnGreen**: pointer to an int in which the green value will be returned
* **pnBlue**: pointer to an int in which the blue value will be returned
* **pnAlpha**: pointer to an int in which the (optional) alpha value will be returned
### Returns
TRUE if the color could be successfully parsed, or FALSE in case of errors.
"""
function getrgbfromstring(hST::Ptr{OGRStyleToolH},pszColor,pnRed,pnGreen,pnBlue,pnAlpha)
    ccall((:OGR_ST_GetRGBFromString,libgdal),Cint,(Ptr{OGRStyleToolH},Cstring,Ptr{Cint},Ptr{Cint},Ptr{Cint},Ptr{Cint}),hST,pszColor,pnRed,pnGreen,pnBlue,pnAlpha)
end

"""
    OGR_STBL_Create(void) -> OGRStyleTableH
OGRStyleTable factory.
### Returns
an handle to the new style table object.
"""
function create()
    checknull(ccall((:OGR_STBL_Create,libgdal),Ptr{OGRStyleTableH},()))
end

"""
    OGR_STBL_Destroy(OGRStyleTableH hSTBL) -> void
Destroy Style Table.
### Parameters
* **hSTBL**: handle to the style table to destroy.
"""
function destroy(hSTBL::Ptr{OGRStyleTableH})
    ccall((:OGR_STBL_Destroy,libgdal),Void,(Ptr{OGRStyleTableH},),hSTBL)
end

"""
    OGR_STBL_AddStyle(OGRStyleTableH hStyleTable,
                      const char * pszName,
                      const char * pszStyleString) -> int
Add a new style in the table.
### Parameters
* **hStyleTable**: handle to the style table.
* **pszName**: the name the style to add.
* **pszStyleString**: the style string to add.
### Returns
TRUE on success, FALSE on error
"""
function addstyle(hStyleTable::Ptr{OGRStyleTableH},pszName,pszStyleString)
    ccall((:OGR_STBL_AddStyle,libgdal),Cint,(Ptr{OGRStyleTableH},Cstring,Cstring),hStyleTable,pszName,pszStyleString)
end

"""
    OGR_STBL_SaveStyleTable(OGRStyleTableH hStyleTable,
                            const char * pszFilename) -> int
Save a style table to a file.
### Parameters
* **hStyleTable**: handle to the style table.
* **pszFilename**: the name of the file to save to.
### Returns
TRUE on success, FALSE on error
"""
function savestyletable(hStyleTable::Ptr{OGRStyleTableH},pszFilename)
    ccall((:OGR_STBL_SaveStyleTable,libgdal),Cint,(Ptr{OGRStyleTableH},Cstring),hStyleTable,pszFilename)
end

"""
    OGR_STBL_LoadStyleTable(OGRStyleTableH hStyleTable,
                            const char * pszFilename) -> int
Load a style table from a file.
### Parameters
* **hStyleTable**: handle to the style table.
* **pszFilename**: the name of the file to load from.
### Returns
TRUE on success, FALSE on error
"""
function loadstyletable(hStyleTable::Ptr{OGRStyleTableH},pszFilename)
    ccall((:OGR_STBL_LoadStyleTable,libgdal),Cint,(Ptr{OGRStyleTableH},Cstring),hStyleTable,pszFilename)
end

"""
    OGR_STBL_Find(OGRStyleTableH hStyleTable,
                  const char * pszName) -> const char *
Get a style string by name.
### Parameters
* **hStyleTable**: handle to the style table.
* **pszName**: the name of the style string to find.
### Returns
the style string matching the name or NULL if not found or error.
"""
function find(hStyleTable::Ptr{OGRStyleTableH},pszName)
    bytestring(ccall((:OGR_STBL_Find,libgdal),Cstring,(Ptr{OGRStyleTableH},Cstring),hStyleTable,pszName))
end

"""
    OGR_STBL_ResetStyleStringReading(OGRStyleTableH hStyleTable) -> void
Reset the next style pointer to 0.
### Parameters
* **hStyleTable**: handle to the style table.
"""
function resetstylestringreading(hStyleTable::Ptr{OGRStyleTableH})
    ccall((:OGR_STBL_ResetStyleStringReading,libgdal),Void,(Ptr{OGRStyleTableH},),hStyleTable)
end

"""
    OGR_STBL_GetNextStyle(OGRStyleTableH hStyleTable) -> const char *
Get the next style string from the table.
### Parameters
* **hStyleTable**: handle to the style table.
### Returns
the next style string or NULL on error.
"""
function getnextstyle(hStyleTable::Ptr{OGRStyleTableH})
    bytestring(ccall((:OGR_STBL_GetNextStyle,libgdal),Cstring,(Ptr{OGRStyleTableH},),hStyleTable))
end

"""
    OGR_STBL_GetLastStyleName(OGRStyleTableH hStyleTable) -> const char *
Get the style name of the last style string fetched with OGR_STBL_GetNextStyle.
### Parameters
* **hStyleTable**: handle to the style table.
### Returns
the Name of the last style string or NULL on error.
"""
function getlaststylename(hStyleTable::Ptr{OGRStyleTableH})
    bytestring(ccall((:OGR_STBL_GetLastStyleName,libgdal),Cstring,(Ptr{OGRStyleTableH},),hStyleTable))
end


text:
function trophicspecies(N::T; clustering_function::Function=EAJS) where {T <: UnipartiteNetwork}

    similarities = clustering_function(N)

    # TODO make a matrix based on overlap

end


text:
module MinimalPerfectHash

export CHD
include("chd.jl")
include("chdhasher.jl")
include("chdconstructor.jl")
#include("chd_new.jl")

end # module

max_stars_repo_name: kyungminlee/MinimalPerfectHashing

max_issues_repo_name: kyungminlee/MinimalPerfectHashing

max_forks_repo_name: kyungminlee/MinimalPerfectHashing


text:
abstract type AbstractTerm{T} end

"""
    getθ(A::AbstractTerm)

Return a vector of the mutable elements of the lower triangle of `A.Λ` (column-major ordering)
"""
function getθ end

getθ(v::Vector{AbstractTerm{T}}) where {T} = vcat(getθ.(v)...)

Lambda(A::AbstractTerm) = Λ(A)     # an alias for when UTF-8 characters are difficult to type

"""
    lowerbd{T}(A::AbstractTerm{T})

Return the vector of lower bounds on the parameters, `θ` associated with `A`

These are the elements in the lower triangle of `A.Λ` in column-major ordering.
Diagonals have a lower bound of `0`.  Off-diagonals have a lower-bound of `-Inf`.
"""
function lowerbd end

"""
    nθ(A::AbstractTerm)

Return the number of free parameters in the relative covariance matrix Λ
"""
function nθ end

"""
    rowlengths(A::AbstractTerm)

Return a vector of the Euclidean lengths of the rows of `A.Λ`
"""
function rowlengths end

"""
    MatrixTerm

Term with an explicit, constant matrix representation

# Members
* `x`: matrix
* `wtx`: weighted matrix
* `cnames`: vector of column names
"""
mutable struct MatrixTerm{T,S<:AbstractMatrix} <: AbstractTerm{T}
    x::S
    wtx::S
    piv::Vector{Int}
    rank::Int
    cnames::Vector{String}
end

function MatrixTerm(X::AbstractMatrix, cnms)
    T = eltype(X)
    cf = cholesky!(Symmetric(X'X, :U), Val(true), tol = -one(T), check = false)
    r = cf.rank
    piv = cf.piv
    X = X[:, piv[1:r]]
    MatrixTerm{T,typeof(X)}(X, X, piv, r, cnms)
end

function MatrixTerm(y::Vector{T}) where {T}
    m = reshape(y, (length(y), 1))
    MatrixTerm{T,Matrix{T}}(m, m, [1], Int(all(iszero, y)), [""])
end

getθ(::MatrixTerm{T}) where {T} = T[]

function getθ!(v::StridedVector{T}, A::MatrixTerm{T}) where T
    @argcheck(length(v) == 0, DimensionMisMatch)
    v
end

lowerbd(::MatrixTerm{T}) where {T} = T[]

nθ(::MatrixTerm) = 0

function reweight!(A::MatrixTerm{T}, sqrtwts::Vector{T}) where {T}
    if !isempty(sqrtwts)
        if (A.x === A.wtx)
            A.wtx = similar(A.x)
        end
        mul!(A.wtx, Diagonal(sqrtwts), A.x)
    end
    A
end

rowlengths(A::MatrixTerm{T}) where {T} = T[]

Base.adjoint(A::AbstractTerm) = Adjoint(A)

Base.eltype(A::MatrixTerm{T}) where {T} = T

Base.length(A::MatrixTerm) = length(A.wtx)

Base.size(A::MatrixTerm) = size(A.wtx)

Base.size(A::Adjoint{T,<:MatrixTerm{T}}) where {T} = reverse(size(A.parent))

Base.size(A::MatrixTerm, i) = size(A.wtx, i)

Base.copyto!(A::MatrixTerm{T}, src::AbstractVecOrMat{T}) where {T} = copyto!(A.x, src)

*(A::Adjoint{T,<:MatrixTerm{T}}, B::MatrixTerm{T}) where {T} = A.parent.wtx'B.wtx

LinearAlgebra.mul!(R::AbstractMatrix{T}, A::MatrixTerm{T}, B::MatrixTerm{T}) where {T} =
    mul!(R, A.wtx, B.wtx)

LinearAlgebra.mul!(R::StridedVecOrMat{T}, A::MatrixTerm{T}, B::StridedVecOrMat{T}) where {T} =
    mul!(R, A.x, B)

LinearAlgebra.mul!(C, A::Adjoint{T,<:MatrixTerm{T}}, B::MatrixTerm{T}) where {T} =
    mul!(C, A.parent.wtx', B.wtx)

Λ(A::MatrixTerm) = I

abstract type AbstractFactorReTerm{T} <: AbstractTerm{T} end

"""
    getθ!{T}(v::AbstractVector{T}, A::AbstractFactorReTerm{T})

Overwrite `v` with the elements of the blocks in the lower triangle of `A.Λ` (column-major ordering)
"""
function getθ! end

"""
    isnested(A::AbstractFactorReTerm, B::AbstractFactorReTerm)

Is factor `A` nested in factor `B`?  That is, does each value of `A` occur with just
one value of B?
"""
function isnested(A::AbstractFactorReTerm, B::AbstractFactorReTerm)
    @argcheck(size(A, 1) == size(B, 1), DimensionMismatch)
    bins = zeros(eltype(B.refs), length(A.levels))
    @inbounds for (a, b) in zip(A.refs, B.refs)
        bba = bins[a]
        if iszero(bba)    # bins[a] not yet set?
            bins[a] = b   # set it
        elseif bba ≠ b    # set to another value?
            return false
        end
    end
    true
end

"""
    levs(A::AbstractFactorReTerm)

Return the levels of the grouping factor.

# Examples
```jldoctest
julia> trm = ScalarFactorReTerm(categorical(repeat('A':'F', inner = 5)), :G);

julia> show(MixedModels.levs(trm))
['A', 'B', 'C', 'D', 'E', 'F']
```
"""
levs(A::AbstractFactorReTerm) = A.levels

"""
    nlevs(A::FactorReTerm)

Return the number of levels in the grouping factor of `A`.
"""
nlevs(A::AbstractFactorReTerm) = length(levs(A))

"""
    nrandomeff(A::FactorReTerm)

Return the total number of random effects in A.
"""
nrandomeff(A::AbstractFactorReTerm) = nlevs(A) * vsize(A)

"""
    vsize(A::AbstractFactorReTerm)

Return the size of vector-valued random effects.
"""
function vsize end

Base.eltype(::AbstractFactorReTerm{T}) where {T} = T

LinearAlgebra.Matrix(A::AbstractFactorReTerm) = Matrix(sparse(A))

LinearAlgebra.mul!(R::Matrix{T}, A::Adjoint{T,<:MatrixTerm{T}},
    B::AbstractFactorReTerm{T}) where {T} = mulαβ!(R, A, B)

*(A::Adjoint{T,<:MatrixTerm{T}}, B::AbstractFactorReTerm{T}) where {T} =
    mul!(Matrix{T}(undef, size(A.parent, 2), size(B, 2)), A, B)

Base.size(A::AbstractFactorReTerm) = (length(A.refs), nrandomeff(A))

Base.size(A::AbstractFactorReTerm, i::Integer) =
    i < 1 ? throw(BoundsError()) : i == 1 ? length(A.refs) :  i == 2 ? nrandomeff(A) : 1

"""
    ScalarFactorReTerm

Scalar random-effects term from a grouping factor

# Members
* `refs`: indices into `levels` for the grouping factor
* `levels`: possible values of the grouping factor
* `z`: the raw random-effects model matrix as a vector.
* `wtz`: a weighted copy of `z`
* `fnm`: the name of the grouping factor as a `Symbol`
* `cnm`: the column name as a string
* `Λ`: the relative covariance multiplier
"""
mutable struct ScalarFactorReTerm{T,R} <: AbstractFactorReTerm{T}
    refs::Vector{R}
    levels::Vector{String}
    z::Vector{T}
    wtz::Vector{T}
    fnm::Symbol
    cnms::Vector{String}
    Λ::T
end
ScalarFactorReTerm(f::CategoricalVector, v::Vector, fnm::Symbol, colnms) =
    ScalarFactorReTerm(CategoricalArrays.order(f.pool)[f.refs], string.(levels(f)), v,
        copy(v), fnm, colnms, one(eltype(v)))
ScalarFactorReTerm(f::CategoricalVector, fnm::Symbol) =
    ScalarFactorReTerm(f, ones(length(f)), fnm, ["(Intercept)"])
ScalarFactorReTerm(f::AbstractVector, fnm::Symbol) = ScalarFactorReTerm(categorical(f),fnm)

getΛ(A::ScalarFactorReTerm) = A.Λ

getθ(A::ScalarFactorReTerm) = [A.Λ]

function getθ!(v::StridedVector, A::ScalarFactorReTerm)
    @argcheck(length(v) == 1, DimensionMismatch)
    v[1] = A.Λ
    v
end

Λ(A::ScalarFactorReTerm) = A.Λ * I

lowerbd(A::ScalarFactorReTerm{T}) where {T} = zeros(T, 1)

nθ(::ScalarFactorReTerm) = 1

function reweight!(A::ScalarFactorReTerm, sqrtwts::Vector)
    if length(sqrtwts) > 0
        if A.z == A.wtz
            A.wtz = A.z .* sqrtwts
        else
            mul!(A.wtz, Diagonal(sqrtwts), A.z)
        end
    end
    A
end

rowlengths(A::ScalarFactorReTerm) = [abs(A.Λ)]

function setθ!(A::ScalarFactorReTerm{T}, v::T) where {T}
    A.Λ = v
    A
end

vsize(A::ScalarFactorReTerm) = 1

function LinearAlgebra.cond(A::ScalarFactorReTerm)
    Λ = A.Λ
    iszero(Λ) ? oftype(Λ, Inf) : one(Λ)
end

function LinearAlgebra.mul!(C::Diagonal{T}, adjA::Adjoint{T,<:ScalarFactorReTerm{T,R}},
        B::ScalarFactorReTerm{T,R}) where {T,R}
    A = adjA.parent
    @argcheck A === B
    d = C.diag
    fill!(d, zero(T))
    @inbounds for (ri, Azi) in zip(A.refs, A.wtz)
        d[ri] += abs2(Azi)
    end
    C
end

function LinearAlgebra.mul!(C::SparseMatrixCSC{T}, adjA::Adjoint{T,<:ScalarFactorReTerm{T}},
        B::ScalarFactorReTerm{T}) where {T}
    A = adjA.parent
    m, n = size(B)
    @argcheck(size(C, 1) == size(A, 2) && n == size(C, 2) && size(A, 1) == m, DimensionMismatch)
    Ar = A.refs
    Br = B.refs
    Az = A.wtz
    Bz = B.wtz
    nz = nonzeros(C)
    rv = rowvals(C)
    fill!(nz, zero(T))
    for k in 1:m       # iterate over rows of A and B
        i = Ar[k]      # [i,j] are Cartesian indices in C - find and verify corresponding position K in rv and nz
        j = Br[k]
        coljlast = Int(C.colptr[j + 1] - 1)
        K = searchsortedfirst(rv, i, Int(C.colptr[j]), coljlast, Base.Order.Forward)
        if K ≤ coljlast && rv[K] == i
            nz[K] += Az[k] * Bz[k]
        else
            throw(ArgumentError("C does not have the nonzero pattern of A'B"))
        end
    end
    C
end

function LinearAlgebra.mul!(C::Matrix{T}, adjA::Adjoint{T,<:ScalarFactorReTerm{T}},
        B::ScalarFactorReTerm{T}) where {T}
    m, n = size(B)
	A = adjA.parent
    @argcheck(size(C, 1)==size(A, 2) && n==size(C, 2) && size(A, 1)==m, DimensionMismatch)
    Ar = A.refs
    Br = B.refs
    Az = A.wtz
    Bz = B.wtz
    fill!(C, zero(T))
    for i in 1:m
        C[Ar[i], Br[i]] += Az[i] * Bz[i]
    end
    C
end

function *(adjA::Adjoint{T,<:ScalarFactorReTerm{T,R}},
        B::ScalarFactorReTerm{T,S}) where {T,R,S}
    A = adjA.parent
    A === B ? mul!(Diagonal(Vector{T}(undef, nlevs(B))), adjA, B) :
    sparse(Vector{Int32}(A.refs), Vector{Int32}(B.refs), A.wtz .* B.wtz)
end

function SparseArrays.sparse(A::ScalarFactorReTerm)
    Az = A.z
    m = length(Az)
    sparse(Vector{Int32}(1:m), Vector{Int32}(A.refs), Az, m, nlevs(A))
end

"""
    VectorFactorReTerm

Random-effects term from a grouping factor, model matrix and block pattern

# Members

* `refs`: indices into `levels` for the grouping factor
* `levels`: possible values of the grouping factor
* `z`: the transposed raw random-effects model matrix
* `wtz`: a weighted copy of `z`
* `wtzv`: a view of `wtz` as a `Vector{SVector{S,T}}`
* `fnm`: the name of the grouping factor as a `Symbol`
* `cnms`: a `Vector` of column names (row names after transposition) of `z`
* `blks`: a `Vector{Int}` of block sizes within `Λ`
* `Λ`: the relative covariance factor
* `inds`: linear indices of θ elements in the relative covariance factor
"""
mutable struct VectorFactorReTerm{T,R,S} <: AbstractFactorReTerm{T}
    refs::Vector{R}
    levels::Vector{String}
    z::Matrix{T}
    wtz::Matrix{T}
    wtzv::Base.ReinterpretArray{SVector{S,T}}
    fnm::Symbol
    cnms::Vector{String}
    blks::Vector{Int}
    Λ::LowerTriangular{T,Matrix{T}}
    inds::Vector{Int}
end
function VectorFactorReTerm(f::CategoricalVector, z::Matrix{T}, fnm, cnms, blks) where {T}
    k, n = size(z)
    @argcheck(k == sum(blks), DimensionMismatch)
    m = reshape(1:abs2(k), (k, k))
    offset = 0
    inds = sizehint!(Int[], (k * (k + 1)) >> 1)
    for kk in blks
        for j in 1:kk, i in j:kk
            push!(inds, m[offset + i, offset + j])
        end
        offset += kk
    end
    VectorFactorReTerm(CategoricalArrays.order(f.pool)[f.refs], string.(levels(f)), z, z,
                       reinterpret(SVector{k,T}, vec(z)), fnm, cnms, blks,
                       LowerTriangular(Matrix{T}(I, k, k)), inds)
end

Λ(A::VectorFactorReTerm{T}) where {T} = RepeatedBlockDiagonal(A.Λ, size(A.z, 2))

function reweight!(A::VectorFactorReTerm{T,R,S}, sqrtwts::Vector) where {T,R,S}
    if !isempty(sqrtwts)
        z = A.z
        if A.wtz === z
            A.wtz = copy(z)
            A.wtzv = reinterpret(SVector{S,T}, vec(A.wtz))
        end
        rmul!(copyto!(A.wtz, z), Diagonal(sqrtwts))
    end
    A
end

function SparseArrays.sparse(A::VectorFactorReTerm{T,R,S}) where {T,R,S}
    n = size(A, 1)
    colind = Matrix{Int32}(undef, S, n)
    rr = A.refs
    @inbounds for j in 1:n
        offset = (rr[j] - 1) * S
        for i in 1:S
            colind[i, j] = offset + i
        end
    end
    sparse(Vector{Int32}(repeat(1:n, inner=S)), vec(colind), vec(A.z))
end

function rowlengths(A::VectorFactorReTerm)
    ld = A.Λ.data
    [norm(view(ld, i, 1:i)) for i in 1:size(ld, 1)]
end

nθ(A::VectorFactorReTerm) = length(A.inds)

getΛ(A::VectorFactorReTerm) = A.Λ

function getθ!(v::StridedVector{T}, A::VectorFactorReTerm{T}) where T
    @argcheck(length(v) == length(A.inds), DimensionMismatch)
    m = A.Λ.data
    @inbounds for (j, ind) in enumerate(A.inds)
        v[j] = m[ind]
    end
    v
end

getθ(A::VectorFactorReTerm) = A.Λ.data[A.inds]

vsize(A::VectorFactorReTerm{T,R,S}) where {T,R,S} = S

lowerbd(A::VectorFactorReTerm{T}) where {T} =
    T[x ∈ diagind(A.Λ.data) ? zero(T) : T(-Inf) for x in A.inds]

lowerbd(v::Vector{AbstractTerm{T}}) where {T} = vcat(lowerbd.(v)...)

function setθ!(A::VectorFactorReTerm{T}, v::AbstractVector{T}) where {T}
    A.Λ.data[A.inds] = v
    A
end

function mulαβ!(C::Matrix{T}, adjA::Adjoint{T,<:MatrixTerm{T}}, B::ScalarFactorReTerm{T,R},
        α=true, β=false) where {T,R}
    A = adjA.parent
    Awt = A.wtx
    n, p = size(Awt)
    m, q = size(B)
    @argcheck(size(C) == (p, q) && m == n, DimensionMismatch)
    isone(β) || rmul!(C, β)
    zz = B.wtz
    @inbounds for (j, rrj) in enumerate(B.refs)
        αzj = α * zz[j]
        for i in 1:p
            C[i, rrj] += αzj * Awt[j, i]
        end
    end
    C
end

function mulαβ!(C::Matrix{T}, adjA::Adjoint{T,<:MatrixTerm{T}},
        B::VectorFactorReTerm{T,R,S}, α=true, β=false) where {T,R,S}
    A = adjA.parent
    Awt = A.wtx
    n, p = size(Awt)
    m, q = size(B)
    @argcheck(size(C) == (p, q) && m == n, DimensionMismatch)
    isone(β) || rmul!(C, β)
    rr = B.refs
    @inbounds for (j,v) in enumerate(B.wtzv)
        coloffset = (rr[j] - 1) * S
        for k in 1:S
            jj = coloffset + k
            avk = α * v[k]
            for i in 1:p
                C[i, jj] += avk * Awt[j, i]
            end
        end
    end
    C
end

LinearAlgebra.cond(A::VectorFactorReTerm) = cond(A.Λ)

LinearAlgebra.mul!(C::Matrix{T}, adjA::Adjoint{T,<:MatrixTerm{T}},
    B::VectorFactorReTerm{T}) where {T} = mulαβ!(C, adjA, B)

function *(adjA::Adjoint{T,<:VectorFactorReTerm{T}}, B::ScalarFactorReTerm{T}) where {T}
    A = adjA.parent
    nzeros = copy(A.wtz)
    k, n = size(nzeros)
    rowind = Matrix{Int32}(undef, k, n)
    refs = A.refs
    bwtz = B.wtz
    for j in 1:n
        bwtzj = bwtz[j]
        offset = k * (refs[j] - 1)
        for i in 1:k
            rowind[i, j] = i + offset
            nzeros[i, j] *= bwtzj
        end
    end
    sparse(vec(rowind), Vector{Int32}(repeat(B.refs, inner=k)), vec(nzeros),
        k * nlevs(A), nlevs(B))
end

*(adjA::Adjoint{T,<:ScalarFactorReTerm{T}}, B::VectorFactorReTerm{T}) where {T} =
    copy(adjoint(B'adjA.parent))

function LinearAlgebra.mul!(C::UniformBlockDiagonal{T},
        adjA::Adjoint{T,<:VectorFactorReTerm{T,R,S}},
        B::VectorFactorReTerm{T,U,P}) where {T,R,S,U,P}
    A = adjA.parent
    @argcheck(A === B)
    Cd = C.data
    @argcheck(size(Cd) == (S, S, nlevs(B)), DimensionMismatch)
    fill!(Cd, zero(T))
    for (r, v) in zip(A.refs, A.wtzv)
        @inbounds for j in 1:S
            vj = v[j]
            for i in 1:S
                Cd[i, j, r] += vj * v[i]
            end
        end
    end
    C
end

function LinearAlgebra.mul!(C::Matrix{T},
        adjA::Adjoint{T,<:VectorFactorReTerm{T,R,S}},
        B::VectorFactorReTerm{T,U,P}) where {T,R,S,U,P}
    A = adjA.parent
    @argcheck((m = size(A, 1)) == size(B, 1) && size(C,1) == size(A,2) &&
        size(C,2) == size(B,2), DimensionMismatch)
    fill!(C, zero(T))
    for (i, j, Av, Bv) in zip(A.refs, B.refs, A.wtzv, B.wtzv)
        ioffset = (i - 1) * S
        joffset = (j - 1) * P
        for (jj, Bvjj) in enumerate(Bv)
            jind = joffset + jj
            for (ii, Avii) in enumerate(Av)
                C[ioffset + ii, jind] += Avii * Bvjj
            end
        end
    end
    C
end

function *(adjA::Adjoint{T,<:VectorFactorReTerm{T,R,S}},
        B::VectorFactorReTerm{T,U,P}) where {T,R,S,U,P}
    A = adjA.parent
    if A === B
        return mul!(UniformBlockDiagonal(Array{T}(undef, S, S, nlevs(B))), adjA, B)
    end
    Az = A.wtzv
    Bz = B.wtzv
    @argcheck((m = size(A, 1)) == size(B, 1), DimensionMismatch)
    ab = S * P
    nz = ab * m
    I = sizehint!(Int32[], nz)
    J = sizehint!(Int32[], nz)
    vals = sizehint!(T[], nz)
    Ar = A.refs
    Br = B.refs
    for i in 1:m
        Azi = Az[i]
        Bzi = Bz[i]
        if iszero(Azi) || iszero(Bzi)
            continue
        end
        Ari = Ar[i]
        Bri = Br[i]
        ioffset = (Ari - 1) * S
        joffset = (Bri - 1) * P
        for jj in 1:P
            jjo = jj + joffset
            Bzijj = Bzi[jj]
            for ii in 1:S
                push!(I, ii + ioffset)
                push!(J, jjo)
                push!(vals, Azi[ii] * Bzijj)
            end
        end
    end
    cscmat = sparse(I, J, vals)
    nzs = nonzeros(cscmat)
    q, r = divrem(length(nzs), S)
    iszero(r) || throw(DimensionMismatch("nnz(cscmat) = $(nnz(cscmat)) should be a multiple of $S"))
    nzasmat = reshape(nzs, (S, q))
    rowblocks = [SubArray{T,1,Vector{T}}[] for i in 1:nlevs(A)]
    rv = rowvals(cscmat)
    inds = 1:S
    pattern = Vector(inds)
    pattern[S] = 0
    for b in 1:q
        rows = view(rv, inds)
        rows .% S == pattern ||
            throw(ArgumentError("Rows for block $b are not contiguous starting at a multiple of $S"))
        push!(rowblocks[div(rows[1], S) + 1], view(nzs, inds))
        inds = inds .+ S
    end
    nlB = nlevs(B)
    colblocks = sizehint!(StridedMatrix{T}[], nlB)
    colrange = 1:P
    for j in 1:nlB
        inds = nzrange(cscmat, colrange[1])
        rows = rv[inds]
        i1 = inds[1]
        for k in 2:P
            inds = nzrange(cscmat, colrange[k])
            rv[inds] == rows || throw(DimensionMismatch("Rows differ ($rows ≠ $(rv[inds])) at column block $j"))
        end
        push!(colblocks, reshape(view(nzs, i1:inds[end]), (length(rows), P)))
        colrange = colrange .+ P
    end
    BlockedSparse(cscmat, nzasmat, rowblocks, colblocks)
end

function LinearAlgebra.mul!(C::Matrix{T}, adjA::Adjoint{T,<:VectorFactorReTerm{T,R,S}},
        B::ScalarFactorReTerm{T}) where {T,R,S}
    m, n = size(B)
    A = adjA.parent
    @argcheck size(C, 1) == size(A, 2) && n == size(C, 2) && size(A, 1) == m DimensionMismatch
    Ar = A.refs
    Br = B.refs
    Az = A.wtz
    Bz = B.wtz
    fill!(C, zero(T))
    for j in 1:m
        offset = S * (Ar[j] - 1)
        Bzj = Bz[j]
        Brj = Br[j]
        for k in 1:S
            C[offset + k, Brj] += Az[k,j] * Bzj
        end
    end
    C
end

function LinearAlgebra.mul!(C::Matrix{T}, adjA::Adjoint{T,<:ScalarFactorReTerm{T}},
	    B::VectorFactorReTerm{T,R,S}) where {T,R,S}
    m, n = size(B)
	A = adjA.parent
    @argcheck size(C, 1) == size(A, 2) && n == size(C, 2) && size(A, 1) == m DimensionMismatch
    Ar = A.refs
    Br = B.refs
    Az = A.wtz
    Bz = B.wtz
    fill!(C, zero(T))
    for j in 1:m
        offset = S * (Br[j] - 1)
        Azj = Az[j]
        Arj = Ar[j]
        for k in 1:S
            C[Arj, offset + k] += Azj * Bz[k,j]
        end
    end
    C
end


text:
# This file is a part of BAT.jl, licensed under the MIT License (MIT).

abstract type AbstractMCMCStats end
AbstractMCMCStats

struct MCMCNullStats <: AbstractMCMCStats end

Base.push!(stats::MCMCNullStats, sv::DensitySampleVector) = stats

Base.append!(stats::MCMCNullStats, sv::DensitySampleVector) = stats

struct MCMCBasicStats{L<:Real,P<:Real} <: AbstractMCMCStats
    param_stats::BasicMvStatistics{P,FrequencyWeights}
    logtf_stats::BasicUvStatistics{L,FrequencyWeights}
    mode::Vector{P}

    function MCMCBasicStats{L,P}(m::Integer) where {L<:Real,P<:Real}
        param_stats = BasicMvStatistics{P,FrequencyWeights}(m)
        logtf_stats = BasicUvStatistics{L,FrequencyWeights}()
        mode = fill(P(NaN), m)

        new{L,P}(
            param_stats,
            logtf_stats,
            mode
        )
    end
end

function MCMCBasicStats(::Type{S}, ndof::Integer) where {
    PT<:Real, T, W, S<:DensitySample{<:AbstractVector{PT},T,W}
}
    SL = promote_type(T, Float64)
    SP = promote_type(PT, W, Float64)
    MCMCBasicStats{SL,SP}(ndof)
end

MCMCBasicStats(chain::MCMCIterator) = MCMCBasicStats(sample_type(chain), totalndof(getdensity(chain)))

function MCMCBasicStats(sv::DensitySampleVector{<:AbstractVector{<:Real}})
    stats = MCMCBasicStats(eltype(sv), innersize(sv.v, 1))
    append!(stats, sv)
end

MCMCBasicStats(sv::DensitySampleVector) = MCMCBasicStats(unshaped.(sv))

function Base.empty!(stats::MCMCBasicStats)
    empty!(stats.param_stats)
    empty!(stats.logtf_stats)
    fill!(stats.mode, eltype(stats.mode)(NaN))

    stats
end

function Base.push!(stats::MCMCBasicStats, s::DensitySample)
    push!(stats.param_stats, s.v, s.weight)
    if s.logd > stats.logtf_stats.maximum
        stats.mode .= s.v
    end
    push!(stats.logtf_stats, s.logd, s.weight)
    stats
end

function Base.append!(stats::MCMCBasicStats, sv::DensitySampleVector)
    for i in eachindex(sv)
        p = sv.v[i]
        w = sv.weight[i]
        l = sv.logd[i]
        push!(stats.param_stats, p, w)  # Memory allocation (view)!
        if sv.logd[i] > stats.logtf_stats.maximum
            stats.mode .= p  # Memory allocation (view)!
        end
        push!(stats.logtf_stats, l, w)
        stats
    end
    stats
end

ValueShapes.totalndof(stats::MCMCBasicStats) = stats.param_stats.m

nsamples(stats::MCMCBasicStats) = stats.param_stats.cov.sum_w

function Base.merge!(target::MCMCBasicStats, others::MCMCBasicStats...)
    for x in others
        if (x.logtf_stats.maximum > target.logtf_stats.maximum)
            target.mode .= x.mode
        end
        merge!(target.param_stats, x.param_stats)
        merge!(target.logtf_stats, x.logtf_stats)
    end
    target
end

Base.merge(a::MCMCBasicStats, bs::MCMCBasicStats...) = merge!(deepcopy(a), bs...)

function reweight_relative!(stats::MCMCBasicStats, reweighting_factor::Real)
    reweight_relative!(stats.param_stats, reweighting_factor)
    reweight_relative!(stats.logtf_stats, reweighting_factor)

    stats
end

function _bat_stats(mcmc_stats::MCMCBasicStats)
    (
        mode = mcmc_stats.mode,
        mean = mcmc_stats.param_stats.mean,
        cov = mcmc_stats.param_stats.cov
    )
end


text:
import Base: <, ==, isless

## PokeRecord

"""
Main record type. May also be used as a pair of Key and Value separately
"""
immutable PokeRecord
    key::Key
    value::Value
end

function Base.show(io::IO, rec::PokeRecord)
    print(io, "PokeRecord($(bytestring(rec.key)))")
end

==(rec1::PokeRecord, rec2::PokeRecord) = rec1.key == rec2.key
<(rec1::PokeRecord, rec2::PokeRecord) = rec1.key < rec2.key
isless(rec1::PokeRecord, rec2::PokeRecord) = rec1 < rec2

## PokeIterator

# TODO: move / create FileStore poke iterator
"""Iterator that reads data from IO one by one"""
type PokeIterator
    io::IO
end

function Base.start(pit::PokeIterator)
    return nothing
end

function Base.next(pit::PokeIterator, s::Void)
    return readobj(pit.io, PokeRecord), nothing
end

function Base.done(pit::PokeIterator, s)
    return eof(pit.io)
end


text:
vis = Visualizer()
if !haskey(ENV, "CI")
    open(vis)
    wait(vis)
end

model = RobotZoo.DoubleIntegrator(2)
TrajOptPlots.set_mesh!(vis, model)
visualize!(vis, model, SA[1,2,0,0.])
TrajOptPlots.set_mesh!(vis, model, color=colorant"pink")
TrajOptPlots.set_mesh!(vis, model, radius=0.5, height=0.5)

delete!(vis)
model = RobotZoo.Pendulum()
TrajOptPlots.set_mesh!(vis, model)
visualize!(vis, model, SA[pi/2,0])
TrajOptPlots.set_mesh!(vis, model, color=colorant"green")
TrajOptPlots.set_mesh!(vis, model, color2=colorant"green")
TrajOptPlots.set_mesh!(vis, model, color=colorant"red", color2=colorant"cyan")

delete!(vis)
model = RobotZoo.Cartpole()
TrajOptPlots.set_mesh!(vis, model)
visualize!(vis, model, SA[-0.5, pi/2, 0,0])

model2 = RobotZoo.Cartpole(0.2,0.5,1.0,9.81)  # make it longer
TrajOptPlots.set_mesh!(vis, model2)
visualize!(vis, model, SA[-0.5, pi/2, 0,0])
TrajOptPlots.set_mesh!(vis, model2, color=colorant"purple")
TrajOptPlots.set_mesh!(vis, model2, color=colorant"purple", color2=colorant"yellow")

delete!(vis)
model = RobotZoo.Acrobot()
TrajOptPlots.set_mesh!(vis, model, color=colorant"purple")
visualize!(vis, model, SA[pi/4, pi/4, 0,0])

model = RobotZoo.Acrobot(l=SA[0.5,1.0])
TrajOptPlots.set_mesh!(vis, model)
visualize!(vis, model, SA[pi/4, pi/4, 0,0])

delete!(vis)
model = RobotZoo.DubinsCar()
TrajOptPlots.set_mesh!(vis, model)
visualize!(vis, model, SA[0.5, 0, pi/4])

model = RobotZoo.DubinsCar(0.5)
TrajOptPlots.set_mesh!(vis, model, height=0.2)
TrajOptPlots.set_mesh!(vis, model, radius=0.1)
TrajOptPlots.set_mesh!(vis, model, color=colorant"cyan")
TrajOptPlots.set_mesh!(vis, model, color=colorant"orange", color2=colorant"gray30")

delete!(vis)
model = RobotZoo.Quadrotor{MRP}()
TrajOptPlots.set_mesh!(vis, model)
g = MRP(expm(SA[1,0,0]*pi/4))
visualize!(vis, model, SA[0,0,1, g.x, g.y, g.z, 0,0,0, 0,0,0])
TrajOptPlots.set_mesh!(vis, model, color=colorant"white")

delete!(vis)
model = RobotZoo.YakPlane{MRP}()
TrajOptPlots.set_mesh!(vis, model)
g = MRP(expm(SA[1,0,0]*(pi/4 + pi)))
visualize!(vis, model, SA[0,0,1, g.x, g.y, g.z, 0,0,0, 0,0,0])
TrajOptPlots.set_mesh!(vis, model, color=colorant"cyan")

import RobotZoo.Satellite
Base.position(::Satellite, x) = @SVector zeros(3)
RobotDynamics.orientation(::Satellite, x) = UnitQuaternion(x[4], x[5], x[6], x[7])

delete!(vis)
model = RobotZoo.Satellite()
TrajOptPlots.set_mesh!(vis, model)
q = rand(UnitQuaternion)
visualize!(vis, model, SA[0,0,0, Rotations.params(q)...])

delete!(vis)
TrajOptPlots.set_mesh!(vis, model, dims=[0.5,0.5,1]*2, color=colorant"red")

# New Rigid Body
struct MyRB{R} <: RobotDynamics.RigidBody{R} end
delete!(vis)
model = MyRB{MRP}()
TrajOptPlots.set_mesh!(vis, model, dims=[0.5,0.5,1]*4, color=colorant"darkblue")

model = RobotZoo.DoubleIntegrator(2)
TrajOptPlots.set_mesh!(vis, model)
visualize!(vis, model, SA[1,2,0,0.])


text:
# Parse the iterator assignment expression from a for loop.
# It is assumed to either have the form "x = <expr>" or "(x, y, z) = <expr>"
# Returns the list of symbols from the left-hand-side and the expression from the right-hand side.
function parse_iterator(it_expr)
    @assert it_expr.head == :(=)
    v_expr = it_expr.args[1]
    fixture = it_expr.args[2]
    if isa(v_expr, Expr)
        @assert v_expr.head == :tuple
        v = v_expr.args
    else
        v = [v_expr]
    end
    v, fixture
end

function build_tuple_unpacks(var_lists)
    vars = []
    unpacks = []
    for var_list in var_lists
        if length(var_list) == 1
            push!(vars, var_list[1])
        else
            varname = gensym("tuple")
            push!(vars, varname)
            escaped_list = map(esc, var_list)
            push!(unpacks, :( ($(escaped_list...),) = $(esc(varname)) ))
        end
    end

    vars, unpacks
end

# Parse the body of a fixture or a testcase.
# This can be either a begin/end block, or a for loop
function parse_body(body_expr)
    if body_expr.head == :for
        iterators = body_expr.args[1]
        body = body_expr.args[2]
        if iterators.head == :block
            pairs = map(parse_iterator, iterators.args)
            var_lists, fixtures = zip(pairs...)
        else
            v, f = parse_iterator(iterators)
            fixtures = [f]
            var_lists = [v]
        end
    elseif body_expr.head == :block
        fixtures = []
        var_lists = []
        body = body_expr
    else
        error("Incorrect body expression of type $(body_expr.head)")
    end

    vars, unpacks = build_tuple_unpacks(var_lists)

    body = esc(body)
    fixtures = map(esc, fixtures)
    vars = map(esc, vars)

    vars, unpacks, fixtures, body
end

# Taken from Test and simplified
function parse_options(options_expr)
    options = :(Dict{Symbol, Any}())
    for arg in options_expr
        if isa(arg, Expr) && arg.head == :(=)
            # we're building up a Dict literal here
            key = Expr(:quote, arg.args[1])
            push!(options.args, Expr(:call, :(=>), key, esc(arg.args[2])))
        else
            error("Unexpected argument $arg to @testcase")
        end
    end
    options
end

"""
    @testcase [option=val ...] <name> begin ... end
    @testcase [option=val ...] <name> for x in fx1, (y, z) in fx2 ... end

Create a testcase object and add it to the current test group.

Available options:

`tags :: Array{Symbol, 1}`: a list of tags for the testcase.
"""
macro testcase(args...)

    options_expr = args[1:end-2]
    name = esc(args[end-1])
    vars, unpacks, fixtures, body = parse_body(args[end])

    if length(options_expr) > 0
        options = parse_options(options_expr)
        tc_call = quote
            Testcase($name, $(fixtures...); $options...) do $(vars...)
                $(unpacks...)
                $body
            end
        end
    else
        tc_call = quote
            Testcase($name, $(fixtures...)) do $(vars...)
                $(unpacks...)
                $body
            end
        end
    end

    :( register_testobj($tc_call) )
end

"""
    @testgroup <name> begin ... end

Create a test group.
The body can contain other [`@testgroup`](@ref) or [`@testcase`](@ref) declarations.
"""
macro testgroup(args...)

    options_expr = args[1:end-2]
    name = esc(args[end-1])
    body = esc(args[end])

    if length(options_expr) > 0
        options = parse_options(options_expr)
        tg_call = quote
            TestGroup($name; $options...) do
                $body
            end
        end
    else
        tg_call = quote
            TestGroup($name) do
                $body
            end
        end
    end

    :( register_testobj($tg_call) )
end

const PRODUCE_VAR = esc(gensym("produce"))

function _fixture(islocal, args...)
    options_expr = args[1:end-1]
    vars, unpacks, fixtures, body = parse_body(args[end])

    fxname = islocal ? :LocalFixture : :GlobalFixture

    if length(options_expr) > 0
        options = parse_options(options_expr)
        quote
            $fxname($(fixtures...); $options...) do $PRODUCE_VAR, $(vars...)
                $(unpacks...)
                $body
            end
        end
    else
        quote
            $fxname($(fixtures...)) do $PRODUCE_VAR, $(vars...)
                $(unpacks...)
                $body
            end
        end
    end
end

"""
    @global_fixture [option=val ...] <name> begin ... end
    @global_fixture [option=val ...] <name> for x in fx1, (y, z) in fx2 ... end

Create a global fixture (a fixture set up once before all
the testcases that use it and torn down after they finish).

The body must contain a single call to [`@produce`](@ref), producing a single value.

The iterables in the `for` loop are either fixtures (constant of global only),
iterable objects or pairs of two iterables used to parametrize the fixture.

Available options:

`instant_teardown :: Bool`: if `true`, the part of the fixture body after the [`@produce`](@ref)
will be executed immediately.

Returns a [`GlobalFixture`](@ref) object.
"""
macro global_fixture(args...)
    _fixture(false, args...)
end

"""
    @local_fixture <name> begin ... end
    @local_fixture <name> for x in fx1, (y, z) in fx2 ... end

Create a local fixture (a fixture set up before each testcase
that uses it and torn down afterwards).

The body must contain a single call to [`@produce`](@ref), producing a single value.

The iterables in the `for` loop are either fixtures (constant of global only),
iterable objects or pairs of two iterables used to parametrize the fixture.

Returns a [`LocalFixture`](@ref) object.
"""
macro local_fixture(args...)
    _fixture(true, args...)
end

"""
    @produce <val> [<label>]

Produce a fixture value (with an optional label).
Must only be called inside the bodies of [`@local_fixture`](@ref) and [`@global_fixture`](@ref).
"""
macro produce(args...)
    args = map(esc, args)
    :( $PRODUCE_VAR($(args...)) )
end


text:
function neumann_series(mat::SparseMatrixCSC)
    inverse = I + mat
    next_term = mat^2

    while nnz(next_term) != 0
        inverse += next_term
        next_term *= mat
    end

    return inverse
end

#=
function make_onelement_array(A)
    isa(A, Array) ? nothing : (A = [A])
    return A
end
 =#

function semvec(observed, imply, loss, diff)

    observed = make_onelement_array(observed)
    imply = make_onelement_array(imply)
    loss = make_onelement_array(loss)
    diff = make_onelement_array(diff)

    #sem_vec = Array{AbstractSem}(undef, maximum(length.([observed, imply, loss, diff])))
    sem_vec = Sem.(observed, imply, loss, diff)

    return sem_vec
end

function get_observed(rowind, data, semobserved;
            args = (),
            kwargs = NamedTuple())
    observed_vec = Vector{semobserved}(undef, length(rowind))
    for i in 1:length(rowind)
        observed_vec[i] = semobserved(
                            args...;
                            data = Matrix(data[rowind[i], :]),
                            kwargs...)
    end
    return observed_vec
end

function skipmissing_mean(mat)
    means = Vector{Float64}(undef, size(mat, 2))
    for i = 1:size(mat, 2)
        @views means[i] = mean(skipmissing(mat[:,i]))
    end
    return means
end

function F_one_person(imp_mean, meandiff, inverse, data, logdet)
    F = logdet
    @. meandiff = data - imp_mean
    F += dot(meandiff, inverse, meandiff)
    return F
end

function remove_all_missing(data)
    keep = Vector{Int64}()
    for i = 1:size(data, 1)
        if any(.!ismissing.(data[i, :]))
            push!(keep, i)
        end
    end
    return data[keep, :], keep
end

function batch_inv!(fun, model)
    for i = 1:size(fun.inverses, 1)
        fun.inverses[i] .= LinearAlgebra.inv!(fun.choleskys[i])
    end
end

#=
function batch_sym_inv_update!(fun::Union{LossFunction, DiffFunction}, model)
    M_inv = inv(fun.choleskys[1])
    for i = 1:size(fun.inverses, 1)
        if size(model.observed.patterns_not[i]) == 0
            fun.inverses[i] .= M_inv
        else
            ind_not = model.observed.patterns_not[i]
            ind = model.observed.patterns[i]

            A = M_inv[ind_not, ind]
            H = cholesky(M_inv[ind_not, ind_not])
            D = H \ A
            out = M_inv[ind, ind] - LinearAlgebra.BLAS.gemm('T', 'N', 1.0, A, D)
            fun.inverses[i] .= out
        end
    end
end =#

function sparse_outer_mul!(C, A, B, ind) #computes A*S*B -> C, where ind gives the entries of S that are 1
    fill!(C, 0.0)
    for i in 1:length(ind)
        BLAS.ger!(1.0, A[:, ind[i][1]], B[ind[i][2], :], C)
    end
end

function sparse_outer_mul!(C, A, ind) #computes A*∇m, where ∇m ind gives the entries of ∇m that are 1
    fill!(C, 0.0)
    @views C .= sum(A[:, ind], dims = 2)
    return C
end

function sparse_outer_mul!(C, A, B::Vector, ind) #computes A*S*B -> C, where ind gives the entries of S that are 1
    fill!(C, 0.0)
    @views @inbounds for i in 1:length(ind)
        C .+= B[ind[i][2]].*A[:, ind[i][1]]
    end
end

function cov_and_mean(rows; corrected = false)
    data = permutedims(hcat(rows...))
    size(rows, 1) > 1 ?
        obs_cov = Statistics.cov(data; corrected = corrected) :
        obs_cov = reshape([0.0],1,1)
    obs_mean = vcat(Statistics.mean(data, dims = 1)...)
    return obs_cov, obs_mean
end

function duplication_matrix(nobs)
    nobs = Int(nobs)
    n1 = Int(nobs*(nobs+1)*0.5)
    n2 = Int(nobs^2)
    Dt = zeros(n1, n2)

    for j in 1:nobs
        for i in j:nobs
            u = zeros(n1)
            u[Int((j-1)*nobs + i-0.5*j*(j-1))] = 1
            T = zeros(nobs, nobs)
            T[j,i] = 1; T[i, j] = 1
            Dt += u*transpose(vec(T))
        end
    end
    D = transpose(Dt)
    return D
end

function elimination_matrix(nobs)
    nobs = Int(nobs)
    n1 = Int(nobs*(nobs+1)*0.5)
    n2 = Int(nobs^2)
    L = zeros(n1, n2)

    for j in 1:nobs
        for i in j:nobs
            u = zeros(n1)
            u[Int((j-1)*nobs + i-0.5*j*(j-1))] = 1
            T = zeros(nobs, nobs)
            T[i, j] = 1
            L += u*transpose(vec(T))
        end
    end
    return L
end

function commutation_matrix(n; tosparse = false)

    M = zeros(n^2, n^2)

    for i = 1:n
        for j = 1:n
            M[i + n*(j - 1), j + n*(i - 1)] = 1.0
        end
    end

    if tosparse M = sparse(M) end

    return M

end

function commutation_matrix_pre_square(A)

    n2 = size(A, 1)
    n = Int(sqrt(n2))

    ind = repeat(1:n, inner = n)
    indadd = (0:(n-1))*n
    for i in 1:n ind[((i-1)*n+1):i*n] .+= indadd end

    A_post = A[ind, :]

    return A_post

end

function commutation_matrix_pre_square_add!(B, A) # comuptes B + KₙA

    n2 = size(A, 1)
    n = Int(sqrt(n2))

    ind = repeat(1:n, inner = n)
    indadd = (0:(n-1))*n
    for i in 1:n ind[((i-1)*n+1):i*n] .+= indadd end

    @views @inbounds B .+= A[ind, :]

    return B

end

function get_commutation_lookup(n2::Int64)

    n = Int(sqrt(n2))
    ind = repeat(1:n, inner = n)
    indadd = (0:(n-1))*n
    for i in 1:n ind[((i-1)*n+1):i*n] .+= indadd end

    lookup = Dict{Int64, Int64}()

    for i in 1:n2
        j = findall(x -> (x == i), ind)[1]
        push!(lookup, i => j)
    end

    return lookup

end

function commutation_matrix_pre_square!(A::SparseMatrixCSC, lookup) # comuptes B + KₙA

    for (i, rowind) in enumerate(A.rowval)
        A.rowval[i] = lookup[rowind]
    end

end

function commutation_matrix_pre_square!(A::SparseMatrixCSC) # computes KₙA
    lookup = get_commutation_lookup(size(A, 2))
    commutation_matrix_pre_square!(A, lookup)
end

function commutation_matrix_pre_square(A::SparseMatrixCSC)
    B = copy(A)
    commutation_matrix_pre_square!(B)
    return B
end

function commutation_matrix_pre_square(A::SparseMatrixCSC, lookup)
    B = copy(A)
    commutation_matrix_pre_square!(B, lookup)
    return B
end

function commutation_matrix_pre_square_add_mt!(B, A) # comuptes B + KₙA # 0 allocations but slower

    n2 = size(A, 1)
    n = Int(sqrt(n2))

    indadd = (0:(n-1))*n

    Threads.@threads for i = 1:n
        for j = 1:n
            row = i + indadd[j]
            @views @inbounds B[row, :] .+= A[row, :]
        end
    end

    return B

end