@model function simulate_data(::Type{T} = Float64;
    DOdata = missing,
    ODE,
    times,
    nDays,
    GPP_daily_mu,
    GPP_daily_sigma,
    ER_daily_mu,
    ER_daily_sigma,
    k600_meanlog,
    k600_sdlog,
    K02_conv,
    err_obs_iid_sigma_scale,
    err_proc_iid_sigma_scale
    ) where {T}

    # Priors on daily GPP and ER, k600, observation error, and process error
    GPP_daily ~ MvNormal(FillArrays.Fill(GPP_daily_mu, nDays), GPP_daily_sigma)
    ER_daily ~ MvNormal(FillArrays.Fill(ER_daily_mu, nDays), ER_daily_sigma)
    k600 ~ LogNormal(k600_meanlog, k600_sdlog)
    kGas = K02_conv*k600
    err_obs_iid_sigma ~ truncated(Cauchy(0, err_obs_iid_sigma_scale), 0, Inf)
    err_proc_iid_sigma ~ truncated(Cauchy(0, err_proc_iid_sigma_scale), 0, Inf)

    p = Dict(
        "PAR" => ODE.p["PAR"],
        "OSat" => ODE.p["OSat"],
        "day" => ODE.p["day"],
        "GPP_daily"=> GPP_daily,
        "ER_daily" => ER_daily,
        "kGas" => kGas,
        "Q" => ODE.p["Q"],
        "V" => ODE.p["V"])

    if DOdata === missing
        DOdata = Vector{T}(undef, length(times))
        DOdata[1] = ODE.u0[1]
    end

    DO_true = Vector{T}(undef, (length(times)))
    DO_true[1] ~ Normal(DOdata[1], err_obs_iid_sigma)

    for i = 2:length(times)
        problem = remake(ODE, p = p, u0 = [DO_true[i-1]], tspan = (times[i-1], times[i]))
        predicted = solve(problem, alg = Trapezoid())
        DO_true[i] ~ Normal(predicted(times[i])[1], err_proc_iid_sigma)
        DOdata[i] ~ Normal(DO_true[i], err_obs_iid_sigma)
    end

    return DOdata
end


using ThreeBodyProblem
using DifferentialEquations
using Plots

μ₁ = 398600 # {km³/s²} gravitational parameter of Earth
μ₂ = 4903   # {km³/s²} gravitational parameter of the Moon
d = 384400  # {km} average distance between Earth and the Moon
p = [μ₁, μ₂, d];

R₁,R₂ = computeR1R2(p)
ωₛ = sqrt((μ₁ + μ₂)/d^3);

Rₑ = 6378.0 # {km} radius of the Earth
Rₘ = 1738.0 # {km} radius of the Moon
h = 200.0   # {km} altitude of parking orbit

L1 = computeL1(p)

vᵢ₀ = 10.92367104   # {km/s} synodic frame injection velocity
ϕ₀ = 47.70061087    # {°}

vᵢ = vᵢ₀
ϕ = ϕ₀

r₀ = [-R₁ - (Rₑ + 200)*cosd(ϕ); -(Rₑ + 200)*sind(ϕ); 0]
v₀ = vᵢ*[sind(ϕ); -cosd(ϕ); 0];
rv₀ = [r₀;v₀]

tspan = [0.,86400*3] # 3 days

prob = ODEProblem(CR3BPdynamics!,rv₀,tspan,p)
sol = solve(prob,reltol=1e-6)

plot(sol,vars=(1,2),title="Initial Guess",label="")
# scatter!(L1[1],L1[2],marker="x")
plot!(circle(Rₑ,[-R₁;0]),color="blue",label="Earth")
plot!(circle(Rₘ,[R₂;0]),color="gray",label="Moon",aspect_ratio=1)


"""
RhoPorosity(args...)

Calculate the porosity from density log

`` { \\phi  }_{ \\rho  }=\\frac { { \\rho  }_{ ma }-{ \\rho  }_{ log } }{ { \\rho  }_{ ma }-{ \\rho  }_{ Fluid } } ``

## Input Example

```julia
RhoPorosity(RhoLog,RhoMa,RhoF)
```
"""
function RhoPorosity(RhoLog,RhoMa,RhoF)
    phi=(RhoMa.-RhoLog)./(RhoMa.-RhoF)
end


######################################################
#
# Canonical Binomial Cascade Function
#
# Date:   2014-09-29
# Author: Hessel Tuinhof
#
# Function to construct a canonical binomial cascade.
# The canonical measure only preserves mass on average.
#
# Reference:
# - Riedi, Rudolf H.: Multifractal Processes
#   (http://www.stat.rice.edu/~riedi/Publ/PDF/MP.pdf)
# - Calvet, L. E.; Fisher, A. J.: Multifractal Volatility
#
# Required packages: Distributions
#
# Input:
#   m0 ... submass, float in (0,1)
#   k ... number of cascades, positive integer
#   startmass ... total mass, default: 1
#
# Output:
#   mu ... vector contains all measures
#
######################################################

using Distributions

function BinomialCanonical(m0, k, startmass = 1)
	mu = zeros(2)
	M = zeros(2)
	for i = 1:2
		rand(Bernoulli()) == 1 ? M[i] = m0 : M[i] = 1-m0
		mu[i] = M[i]*startmass
	end
	if k > 1
		tmp = Float64[]
		mu = append!(tmp, [BinomialCanonical(m0, k-1, mu[1]), BinomialCanonical(m0, k-1, mu[2])])
	end
	return mu
end


# datagen.jl -- Generates CSV file containing 4 rows:
#               Row 1 contains resistor values from 0 to 100k Ohms
#               Row 2 contains capacitor values from 0 to 1u Farad
#               Row 3 contains times from 0 to 5RC (elementwise)
#               Row 4 contains outputs
#
# Accepts 2 inputs:
#               Input 1 is the number of datpoints to generate.
#               Input 2 is the filename to save the results to.

using Random
using DelimitedFiles

numRows = parse(UInt64, ARGS[1]);
vInitial = 10; # Value at t = 0

data = rand!(zeros(4,numRows));
data[1,:] *= 100000; # Put resistor column at 10k order of magnitude
data[2,:] /= 100000; # Put capacitor column at 1u order of magnitude

for i = 1:numRows
        data[3,i] *= 5*data[1,i]*data[2,i]; # Keeps t less than 5RC at each point
end

for i = 1:numRows
        data[4,i] = vInitial * exp(-data[3,i] / (data[1,i]*data[2,i])); # Ckt output value
end

writedlm(ARGS[2], data, ','); # Creates CSV file for data array.


include("../matplotlib_settings.jl")

using GeometricExamples

include("guiding_center_4d_settings_vprk.jl")

tableau_list = get_tableau_list_vprk_projection()

equ = problem.guiding_center_4d_iode()

include("guiding_center_4d.jl")

run_tableau_list(tableau_list, equ, Δt, ntime)

generate_html(tableau_list, figures, prefix, fig_suff, fig_dir, out_dir, html_file;
              title=coll_title, params=get_params(Δt, ntime, problem))


function inverse_variance(μ::AbstractVector, Σ::M) where {M<:Union{AbstractMatrix,AbstractPDMat}}
    n = size(μ)
    denom = sum(Σ \ ones(n))
    τhat = sum(Σ \ μ) / denom
    Vτhat = 1.0/denom
    τpost=Normal(τhat, √Vτhat)
    return τpost
end

function unweighted_mean(μ::AbstractVector, Σ::AbstractMatrix)
    n = length(μ)
    τhat = mean(μ)
    Vτhat = sum(Σ) / n^2
    τpost=Normal(τhat, √Vτhat)
    return τpost
end
function unweighted_mean(μ::AbstractVector, Σ::AbstractPDMat)
    return unweighted_mean(μ, Matrix(Σ))
end

function weighted_mean(μ::AbstractVector, Σ::M, w::AbstractVector) where {M<:Union{AbstractMatrix,AbstractPDMat}}
    τhat = sum(μ .* w) / sum(w)
    Vτhat = dot(w, Σ*w) / sum(w)^2
    τpost=Normal(τhat, √Vτhat)
    return τpost
end
# function weighted_mean(μ::AbstractVector, Σ::AbstractPDMat, w)
    # return weighted_mean(μ, PDMats.full(Σ), w)
# end


"""Generates *n* points at random on a *N*-dimensional ball."""
function on_ball(N::Int, n::Int)
    p = rand(Normal(), N, n)
    r = sqrt.(sum(abs2, p, dims=1))
    return p./r
end

"""Generates *n* points at random in a *N*-dimensional ball."""
function in_ball(N::Int, n::Int)
    X = reshape(rand(Normal(), n*N), N, n)
    Y = rand(Poisson(), n)
    r = sqrt.(Y' .+ sum(abs2, X, dims=1))
    return X./r
end

"""Generates *m* random *N*-dimensional translation vectors with *μ_bounds* bounds"""
function translations(N::Int, n::Int, μ_bounds::Tuple{Float64,Float64})
    o_min, o_max = μ_bounds
    Dτ = Uniform(o_min, o_max)
    o = on_ball(N, n)
    τ = rand(Dτ, n)
    return (o.*τ' .+ fill(1.0,N,n))./2.
end

"""Generates manifold basis and its complement"""
function lm_basis(N::Int, M::Int)
    @assert N > M
    X = rand(Normal(),N,N+1)
    Y = (X.-mean(X, dims=2))'./sqrt(N)
    F = svd(Y)
    F.V[:,1:M], F.V[:,(M+1):end]
end

"""Generates *n* points of a *M*-dimensional linear manifold cluster
situated in a hypercube."""
function generate_cluster(
                  n::Int,                   # number of points in generated manifold
                  μ::Vector{Float64},       # manifold translation vector
                  B::Matrix{Float64},       # manifold basis
                  B′::Matrix{Float64},      # orthogonal complement to manifold basis
                  DΦ::Vector{T},  # bound of manifold points
                  DE::Vector{T}   # bound of a point extent from a manifold
                ) where T<:Distribution
    N, M = size(B)
    @assert size(B′) == (N, N-M) "NULL space dimention must be $(N-M)"
    @assert length(DΦ) == M
    @assert length(DE) == N-M

    mfld = zeros(N,n)
    c = 1
    while c <= n
        ϕ = map(d->rand(d), DΦ)
        ɛ = map(d->rand(d), DE)
        x = μ + B*ϕ + B′*ɛ
        # check if manifold point outside of unit hypercube
        b = extrema(x)
        if b[1] < -1 || b[2] > 1.
            continue
        end
        mfld[:,c] = x
        c+=1
    end
    return mfld
end

"""Generates *n* points of a *M*-dimensional linear manifold cluster
situated in a hypercube."""
function generate_cluster(
                  n::Int,                   # number of points in generated manifold
                  μ::Vector{Float64},       # manifold translation vector
                  B::Matrix{Float64},       # manifold basis
                  B′::Matrix{Float64},      # orthogonal complement to manifold basis
                  Φ::Vector{Float64},       # bound of manifold points
                  E::Vector{Float64}        # bound of a point extent from a manifold
                )
    DΦ = map(ϕ->Uniform(-ϕ, ϕ), Φ)
    DE = map(ɛ->Uniform(-ɛ, ɛ), E)
    return generate_cluster(n, μ, B, B′, DΦ, DE)
end

"""Generates *n* points of a *M*-dimensional linear manifold cluster
situated in a hypercube."""
function generate_cluster(
                  n::Int,                   # number of points in generated manifold
                  μ::Vector{Float64},       # manifold translation vector
                  B::Matrix{Float64},       # manifold basis
                  B′::Matrix{Float64},      # orthogonal complement to manifold basis
                  Φ::Float64,               # bound of manifold points
                  E::Float64,               # bound of a point extent from a manifold
                )
    N, M = size(B)
    DΦ = fill(Uniform(-Φ, Φ), M)
    DE = fill(Uniform(-E, E), N-M)
    return generate_cluster(n, μ, B, B′, DΦ, DE)
end

"""Generates *n* points of a *M*-dimensional linear manifold cluster
situated in a hypercube using normal distribution."""
function generate_cluster(
                  ::Type{Normal},
                  n::Int,                   # number of points in generated manifold
                  μ::Vector{Float64},       # manifold translation vector
                  B::Matrix{Float64},       # full basis
                  M::Int,                   # cluster dimension
                  Φ::Float64,               # bound of manifold points
                  E::Float64,               # bound of a point extent from a manifold
                )
    N = size(B,1)
    DΦ = fill(Normal(0, Φ), M)
    DE = fill(Normal(0, E), N-M)
    return generate_cluster(n, μ, B[:,1:M], B[:, M+1:end], DΦ, DE)
end

"""Generates *n* points of a *M*-dimensional linear manifold cluster
situated in a hypercube using uniform distribution."""
function generate_cluster(
                  ::Type{Uniform},
                  n::Int,                   # number of points in generated manifold
                  μ::Vector{Float64},       # manifold translation vector
                  B::Matrix{Float64},       # full basis
                  M::Int,                   # cluster dimension
                  Φ::Float64,               # bound of manifold points
                  E::Float64,               # bound of a point extent from a manifold
                )
    N = size(B,1)
    DΦ = fill(Uniform(-Φ, Φ), M)
    DE = fill(Uniform(-E, E), N-M)
    return generate_cluster(n, μ, B[:,1:M], B[:, M+1:end], DΦ, DE)
end

""" Generates *m* linear manifold clusters
where *m* is size of parameter *M* """
function generate(n::Int,                     # number of points in generated manifold
                  N::Int,                     # space dimensionality
                  M::Vector{Int},             # generated manifolds dimensions
                  τ::Tuple{Float64,Float64},  # translation vector bounds
                  Φ::Vector{Vector{Float64}}, # bound of manifold points
                  E::Vector{Vector{Float64}}  # bound of a point extent from a manifold
                )
    m = length(M) # number of manifolds to generate

    @assert ( length(E) == m && length(Φ) == m ) "Bounds should correspond to number of generated manifolds"

    μ = translations(N, m, τ)
    bases = Matrix{Float64}[]
    manifolds = Matrix{Float64}[]
    for i in 1:m
        B, B′ = lm_basis(N, M[i])
        mfld = generate_cluster(n, μ[:,i], B, B′, Φ[i], E[i])
        push!(bases, B)
        push!(manifolds, mfld)
    end
    return manifolds, bases, μ
end

"""Generate manifold dimensions."""
function mdims(N, m, dim_cut::Int = 7)
    M = cumsum(ones(Int,m))
    if m >= N
        n = N >= dim_cut ? dim_cut : N
        M[n:end] .= n
    end
    return M
end

""" Generates *m* linear manifold clusters
where *m* is size of parameter *M* """
function generate(n::Int,                     # number of points in generated manifold
                  N::Int,                     # space dimensionality
                  M::Vector{Int},             # number of manifolds
                  τ::Tuple{Float64,Float64},   # translation vector bounds
                  κ::Float64
                )
    m = length(M)

    # generate bounds
    D = Uniform(0.2, 0.99); rand(D, 1000);
    ΔΦ = [sort(rand(D, 2))::Vector{Float64} for i in 1:m]
    ΔE = [sort(rand(D, 2)*diff(mean(ΔΦ))[1]/κ)::Vector{Float64} for i in 1:m]

    Φ = Vector{Float64}[]
    E = Vector{Float64}[]
    for i in M
        DΦ = Uniform(ΔΦ[i][1], ΔΦ[i][2]); rand(DΦ, 1000);
        DE = Uniform(ΔE[i][1], ΔE[i][2]); rand(DE, 1000);
        push!(Φ, rand(DΦ, M[i]))
        push!(E, rand(DE, N-M[i]))
    end
    return generate(n, N, M, τ, Φ, E)
end

""" Generates *m* linear manifold clusters"""
function generate(n::Int,                     # number of points in generated manifold
                  N::Int,                     # space dimensionality
                  m::Int,                     # number of manifolds
                  τ::Pair{Float64,Float64},   # translation vector bounds
                  κ::Float64;
                  dim_cut::Int = 7
                )
    # generate manifold dimensions
    M = mdims(N-1,m,dim_cut)
    return generate(n, N, M, (τ.first,τ.second), κ)
end

function noise(N::Int, n::Int)
    return rand(N, n)
end


function generate_integrand_uv(kernel, testref, trialref, testel, trialel)

    function k3(u,v)
        out = @SMatrix zeros(3,3)

        x = neighborhood(testel,u)
        y = neighborhood(trialel,v)

        kernelval = kernel(x,y)
        f = testref(x)
        g = trialref(y)

        jx = jacobian(x)
        jy = jacobian(y)
        ds = jx*jy

        return  SMatrix{3,3}([dot(f[i][1], kernelval*g[j][1])*ds for i=1:3, j=1:3])
    end

    return k3
end


# Create example data
x = LinRange(-2, 2, 10)
y = x.^3 .+ x.^2 .+ x .+ 6
err = LinRange(0.5, 3, 10)
# Draw symmetric, horizontal error bars
errorbar(x, y, err, horizontal=true)
# Draw asymmetric error bars with markers
errorbar(x, y, err, err./2, "-o")


include("header.jl")

#=
# Fluid statics
In this notebook, we will learn
* About the specific weight of a fluid and how pressure increases with depth
* The concept of a fluid interface and a free surface
* How to perform basic calculations
=#

# ### Set up the module
using MAE103
#-
using Plots

#=
Fluid statics basically involves fluids that are not in motion (although we'll generalize
this later) It is described by a basic equation that governs the pressure,

$$\dfrac{\mathrm{d}p}{\mathrm{d}z} = -\rho g$$

where $\rho$ is the density of the fluid, $g$ is gravity, and $z$ is a coordinate
pointing vertically upward (i.e., $g$ is pointing in the $-z$ direction). The grouping
$\rho g$ is called the *specific weight* of the fluid, and is usually given the symbol
$\gamma$.

For example, for water, compare
=#
ρg = Density(Water)*Gravity()
# with
γ = SpecificWeight(Water)

#=
This equation shows us that pressure increases with depth, in the direction of gravity.
However, it is important to understand that it only tells us about *differences*
of pressure, not absolute pressure:

$$p_2 - p_1 = -\gamma (z_2 - z_2)$$

Notice that it doesn't matter how wide the fluid is. It could be in the ocean
or in a small tube. Only depth matters, since it is the weight of fluid *per unit
area* above a certain position that contributes to the pressure.
=#
#-
#=
To give an example of how we use this, let's first define the concept of a *free
surface*: the interface between a liquid and a gas. The most obvious example of
this is where liquid and air meet, like at the surface of the ocean or
in a bath of water. Gases are much less dense than liquids at standard conditions,
so they have much smaller specific weights.
=#
SpecificWeight(Air)
#=
This means that the pressure
variations in a static gas are much smaller than in a static liquid. As a result, it is safe to assume that (absolute) pressure in the gas is uniform, and
equal to some ambient value. For example, in air, we often assume it is 1 atm.
We often use the symbol $p_\infty$ for this ambient pressure:
=#
p∞ = Pressure(1u"atm")

#=
We can obtain specific weights in liquids other than water by knowing their
*specific gravity*, equal to their specific weight divided by that of water.
For example, if a certain type of oil has specific gravity of 0.9, it means its
density is 0.9 times that of water:
=#
γoil = SpecificWeight(SpecificGravity(0.9))

#=
So let's get an appreciation for how much the pressure increases in water due to fluid
statics. For example, in 1 m of water, the pressure increases by
=#
dp = Pressure(SpecificWeight(Water)*1u"m")

#=
How much is this increase in atmospheres?
=#
value(dp,u"atm")
#=
so for every meter of depth, the pressure increases by 10 percent of one
atmosphere. This means we only need to go down to 10 meters (or 34 feet)
to increase the pressure by one atmosphere. That's not far at all.

We can plot the pressure distribution through this depth by using `staticpressure`
=#
staticpressure(10u"m",Water)

#-
#=
Because the fluid statics equation describes pressure differences, then it can be
used in a *cumulative* manner. In other words, to calculate the pressure at some
point, we add the contributions of all of the fluid (or fluids) above that point.

For example, to emphasize how little the pressure in air varies with depth,
let's imagine there is a 10 m layer of air on top of the 10 m layer of water,
and plot the pressure distribution throughout both layers. The top of the air
layer is at 1 atm. (Here, we are imagining both fluids sitting inside a container,
so a layer of air makes sense.)
=#
staticpressure([10u"m",20u"m"],[Air,Water])

#=
There's no noticeable change of pressure in the upper (air) layer, so we
can always assume that the pressure is uniform in static air. But the pressure in
the water increases much more rapidly with depth.
=#
#-
#=
### Example
A 10-foot layer of glycerin lies below a 5-ft layer of water. The water is exposed
to pressurized air at $p_\infty = 3$ atm. Determine the pressure at the bottom of the glycerin
layer, in Pa.

We simply add the contributions in each layer, and the answer will be

$p = p_\infty + \gamma_w h_w + \gamma_g h_g$

where $\gamma_w$ is the specific weight of water, $h_g$ is the height of the
water layer, $\gamma_g$ the specific weight of glycerin, and $h_g$ the
height of this layer. As usual, we need to make sure that the units are consistent
with each other when we perform this calculation. For example, to get
the height of the water layer in meters,
=#
h_w = Height(5u"ft")

#=
But actually the tools in our notebooks do this all for us:
=#
Pressure(3u"atm" + SpecificWeight(Water)*5u"ft" + SpecificWeight(Glycerin)*10u"ft")

#=
We can plot the pressure in these layers:
=#
staticpressure([5u"ft",5u"ft"+10u"ft"],[Water,Glycerin],ambient=3u"atm")

#=
Glycerin is only slightly denser than water, so it doesn't look remarkably
different in the two layers.
=#


# test pd matrix types
using PDMats
using Base.Test

call_test_pdmat(p::AbstractPDMat,m::Matrix) = test_pdmat(p,m,cmat_eq=true,verbose=1)

for T in [Float64,Float32]
  #test that all external constructors are accessible
  m = eye(T,2) ; @test PDMat(m,cholfact(m)).mat == PDMat(Symmetric(m)).mat == PDMat(m).mat == PDMat(cholfact(m)).mat
  d = ones(T,2) ; @test PDiagMat(d,d).inv_diag == PDiagMat(d).inv_diag
  x = one(T) ; @test ScalMat(2,x,x).inv_value == ScalMat(2,x).inv_value
  #for Julia v0.3, do not test sparse matrices with Float32 (see issue #14076)
  (T == Float64 || VERSION >= v"0.4.2") && (s = speye(T,2,2) ; @test PDSparseMat(s,cholfact(s)).mat == PDSparseMat(s).mat == PDSparseMat(cholfact(s)).mat)

  #test the functionality
  M = convert(Array{T,2},[4. -2. -1.; -2. 5. -1.; -1. -1. 6.])
  V = convert(Array{T,1},[1.5, 2.5, 2.0])
  X = convert(T,2.0)

  call_test_pdmat(PDMat(M),M) #tests of PDMat
  call_test_pdmat(PDiagMat(V),diagm(V)) #tests of PDiagMat
  call_test_pdmat(ScalMat(3,x),x*eye(T,3)) #tests of ScalMat
  #for Julia v0.3, do not test sparse matrices with Float32 (see issue #14076)
  (T == Float64 || VERSION >= v"0.4.2") && call_test_pdmat(PDSparseMat(sparse(M)),M)
end

m = eye(Float32,2)
@test convert(PDMat{Float64}, PDMat(m)).mat == PDMat(convert(Array{Float64}, m)).mat
m = ones(Float32,2)
@test convert(PDiagMat{Float64}, PDiagMat(m)).diag == PDiagMat(convert(Array{Float64}, m)).diag
x = one(Float32); d = 4
@test convert(ScalMat{Float64}, ScalMat(d, x)).value == ScalMat(d, convert(Float64, x)).value
if VERSION >= v"0.4.2"
    s = speye(Float32, 2, 2)
    @test convert(PDSparseMat{Float64}, PDSparseMat(s)).mat == PDSparseMat(convert(SparseMatrixCSC{Float64}, s)).mat
end


# ---
# title: 715. Range Module
# id: problem715
# author: Tian Jun
# date: 2020-10-31
# difficulty: Hard
# categories: Segment Tree, Ordered Map
# link: <https://leetcode.com/problems/range-module/description/>
# hidden: true
# ---
#
# A Range Module is a module that tracks ranges of numbers. Your task is to
# design and implement the following interfaces in an efficient manner.
#
# * `addRange(int left, int right)` Adds the half-open interval `[left, right)`, tracking every real number in that interval. Adding an interval that partially overlaps with currently tracked numbers should add any numbers in the interval `[left, right)` that are not already tracked.
#
# * `queryRange(int left, int right)` Returns true if and only if every real number in the interval `[left, right)` is currently being tracked.
#
# * `removeRange(int left, int right)` Stops tracking every real number currently being tracked in the interval `[left, right)`.
#
# **Example 1:**
#
#
#
#     **addRange(10, 20)** : null
#     **removeRange(14, 16)** : null
#     **queryRange(10, 14)** : true (Every number in [10, 14) is being tracked)
#     **queryRange(13, 15)** : false (Numbers like 14, 14.03, 14.17 in [13, 15) are not being tracked)
#     **queryRange(16, 17)** : true (The number 16 in [16, 17) is still being tracked, despite the remove operation)
#
#
# **Note:**
#
# * A half open interval `[left, right)` denotes all real numbers `left <= x < right`.
# * `0 < left < right < 10^9` in all calls to `addRange, queryRange, removeRange`.
# * The total number of calls to `addRange` in a single test case is at most `1000`.
# * The total number of calls to `queryRange` in a single test case is at most `5000`.
# * The total number of calls to `removeRange` in a single test case is at most `1000`.
#
#
## @lc code=start
using LeetCode

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


using LinearAlgebraicRepresentation
Lar = LinearAlgebraicRepresentation
using ViewerGL
GL = ViewerGL

n = 200 #3000 #1000 #20000
t = 0.35 #0.15 #0.4 #0.15

function randlines(n=1000,t=0.2)
	V = zeros(Float64,2,2*n)
	EV = [zeros(Int64,2) for k=1:n]
	for k=1:n
		v1 = rand(Float64,2)
		v2 = rand(Float64,2)
		vm = (v1+v2)/2
		transl = rand(Float64,2)
		V[:,k] = (v1-vm)*t + transl
		V[:,n+k] = (v2-vm)*t + transl
		EV[k] = [k,n+k]
	end
	V = GL.normalize2(V)
	model = (V,EV)
	return model
end

model = randlines(n,t)
Sigma = Lar.spaceindex(model)
GL.VIEW([ GL.GLLines(model...)] )

W,EW = Lar.fraglines(1.5,1.5,1.5)(model)
GL.VIEW([ GL.GLLines(W,EW,GL.COLORS[1])] );

W,EW = Lar.fragmentlines(model)
V,EVs = Lar.biconnectedComponent((W,EW::Lar.Cells)) # 2-connected components (H & T)

graphs = [ GL.GLLines(V,EVs[k],GL.COLORS[(k-1)%12+1]) for k=1:length(EVs) ]
GL.VIEW(graphs);


# function combinevectors(x::Vector{Vector{T}})::Vector{T} where T

#     if isempty(x)
#         return Vector{T}(undef, 0)
#     end

#     N = sum(length(x[i]) for i = 1:length(x))

#     y = Vector{T}(undef,N)

#     st_ind = 0
#     fin_ind = 0
#     for i = 1:length(x)
#         st_ind = fin_ind + 1
#         fin_ind = st_ind + length(x[i]) - 1

#         y[st_ind:fin_ind] = x[i]
#     end

#     return y
# end

function array2matrix(X::Vector{Vector{T}})::Matrix{T} where T

    N = length(X)
    D = length(X[1])

    out = Matrix{T}(undef,D,N)
    for n = 1:N
        out[:,n] = X[n]
    end

    return out
end

###########  density fit-related

# Boyd's convex optimization convention: log(-2) := Inf.
function evalcvxlog(x::T)::T where T <: Real
    #println("x = ", x)
    if x < zero(T)
        return Inf
    end

    return log(x)
end

function RKHSfitdensitycostfunc(α::Vector{T}, Kp::Matrix{T}, y::Vector{T}, μ::T)::T where T
    N = length(y)
    @assert length(α) == size(Kp,1) == size(Kp,2)

    r = Kp*α-y
    term1 = dot(r,r)
    term2 = μ*dot(α, Kp, α)
    obj = term1 + term2

    return obj
end
# # test code.
# θ = θ_canonical
# y = vec(𝑌)
# μ = σ²
#
# K = constructkernelmatrix(𝑋, θ)
# #α_solution = fitpdfviaSDP(y, K, μ, max_iters)
#
# # verify objective.
# obj = mycostfunc(α_SDP, K, y, μ)
# # this objective scaoe is the same as the cone solver's. 1.69e-01 answer.

########### end of density fit-related

function findboundingbox(X::Vector{Vector{T}}) where T
    D = length(X[1])
    N = length(X)

    limit_a = Vector{T}(undef, D)
    limit_b = Vector{T}(undef, D)

    for d = 1:D
        Xd = collect( X[n][d] for n = 1:N )
        sorted_Xd = sort(Xd)

        limit_a[d] = sorted_Xd[1]
        limit_b[d] = sorted_Xd[end]
    end

    return limit_a, limit_b
end

function bandpasstotalvariations(LP_percentages,
                        LP_attenuations,
                        HP_percentages,
                        HP_attenuations,
                        y::Vector{T}) where T
    #
    N = length(LP_percentages)
    @assert N == length(LP_attenuations) ==
                length(HP_percentages) ==
                length(HP_attenuations)

    #
    tv_array = Vector{T}(undef, N)

    for i = 1:N
        y_bp = SignalTools.applybandpass( LP_percentages[i],
                                LP_attenuations[i],
                                HP_percentages[i],
                                HP_attenuations[i],
                                y )
        #
        tv_array[i] = sum( abs.(y_bp) )
    end

    return tv_array
end

function bandpassmaxmagnitudes(LP_percentages,
                        LP_attenuations,
                        HP_percentages,
                        HP_attenuations,
                        y::Vector{T}) where T
    #
    N = length(LP_percentages)
    @assert N == length(LP_attenuations) ==
                length(HP_percentages) ==
                length(HP_attenuations)

    #
    max_magnitude = Vector{T}(undef, N)

    for i = 1:N
        y_bp = SignalTools.applybandpass( LP_percentages[i],
                                LP_attenuations[i],
                                HP_percentages[i],
                                HP_attenuations[i],
                                y )
        #
        max_magnitude[i] = maximum( abs.(fft(y_bp)) )
    end

    return max_magnitude
end

# 1. FGPGM for parameter identification in systems of nonlinear ODEs,
#       supplementary.
# Let V be the derivative variables of X.
# D_mat is equation 34 of [1]
function evalderivative(x::Vector{T},
                        μ_GP::Vector{T},
                        μ_GP_dot::Vector{T},
                        θ::KT) where {T <: Real, KT}

    #
    K_XX, K_XV, K_VV = constructkernelematrices(X, θ)

    #K_VX = copy(K_XV')
    μ_dot, S_dot = Utilities.getMVNmarginalparams(x, μ_GP_dot, μ_GP,
                                    K_VV, K_XV, K_XX)
    #
    D_mat = K_XV*inv(K_XX)

    return μ_dot, S_dot, D_mat
end

# b splines

# constant (flat) boundary condition.
# see https://github.com/JuliaMath/Interpolations.jl/blob/master/src/b-splines/cubic.jl
function cubicitp1D(c, x::T, a::T, b::T) where T <: Real
    if one(T) <= x <= length(c)

        δx = abs(x - floor(x))
        i = round(Int, x-δx)

        # out = 1/6*c[i-1]*(1+i-δx)^3 + c[i]*(2/3 - (δx-i)^2 + 1/2*(δx-i)^3) +
        #          c[i+1]*(2/3 - (1+i-δx)^2 + 1/2*(1+i-δx)^3) + 1/6*c[i+2]*(δx-i)^3

        out = 1/6*c[i-1]*(1+i-x)^3 + c[i]*(2/3 - (x-i)^2 + 1/2*(x-i)^3) +
                 c[i+1]*(2/3 - (1+i-x)^2 + 1/2*(1+i-x)^3) + 1/6*c[i+2]*(x-i)^3

        return out
    end

    if x < one(T)
        return a
    end

    return b
end

# # ignore boundary conditions for now.
function cubicitp2D(c, x::Vector{T}, N::Vector{Int}, a, b) where T <: Real

    @assert one(T) <= x[1] <= convert(T,N[1])
    @assert one(T) <= x[2] <= convert(T,N[2])

    δx1 = abs(x[1] - floor(x[1]))
    i = round(Int, x[1]-δx1)

    δx2 = abs(x[2] - floor(x[2]))
    j = round(Int, x[2]-δx2)

    running_sum = zero(T)
    for 𝑖 in Iterators.product(i-1:i+2, j-1:j+2)
        𝑖_vec = collect(𝑖)
        running_sum += c[𝑖_vec...]*evalBspline3(x-𝑖_vec)
    end

    return running_sum
end

function evalBspline3(x::Vector{T}) where T <: Real
    return prod( evalBspline3(x[d]) for d = 1:length(x) )
end

function evalBspline3(x::T) where T <: Real
    abs_x = abs(x)

    if zero(T) <= abs_x < one(T)
        return 2/3 - abs_x^2 + 0.5*abs_x^3
    end

    if one(T) <= abs_x < 2.0
        return 1/6 * (2-abs_x)^3
    end

    return zero(T)
end

function cubicitp2Dexplicit(c, x::Vector{T}, N::Vector{Int}) where T <: Real

    @assert one(T) <= x[1] <= convert(T,N[1])
    @assert one(T) <= x[2] <= convert(T,N[2])

    # determine anchor, 𝑖 = (i1,i2)
    δx1 = abs(x[1] - floor(x[1]))
    i1 = round(Int, x[1]-δx1)

    δx2 = abs(x[2] - floor(x[2]))
    i2 = round(Int, x[2]-δx2)

    x1_lag = x[1]-i1
    x2_lag = x[2]-i2

    x1_adv = 1 - x1_lag
    x2_adv = 1 - x2_lag

    out  = c[i1-1, i2-1]*evalBspline3far(x1_lag)*evalBspline3far(x2_lag)
    out += c[i1, i2-1]*evalBspline3near(x1_lag)*evalBspline3far(x2_lag)
    out += c[i1-1, i2]*evalBspline3far(x1_lag)*evalBspline3near(x2_lag)
    out += c[i1, i2]*evalBspline3near(x1_lag)*evalBspline3near(x2_lag)

    out += c[i1-1, i2+1]*evalBspline3far(x1_lag)*evalBspline3near(x2_adv)
    out += c[i1, i2+1]*evalBspline3near(x1_lag)*evalBspline3near(x2_adv)
    out += c[i1-1, i2+2]*evalBspline3far(x1_lag)*evalBspline3far(x2_adv)
    out += c[i1, i2+2]*evalBspline3near(x1_lag)*evalBspline3far(x2_adv)

    out += c[i1+1, i2-1]*evalBspline3near(x1_adv)*evalBspline3far(x2_lag)
    out += c[i1+2, i2-1]*evalBspline3far(x1_adv)*evalBspline3far(x2_lag)
    out += c[i1+1, i2]*evalBspline3near(x1_adv)*evalBspline3near(x2_lag)
    out += c[i1+2, i2]*evalBspline3far(x1_adv)*evalBspline3near(x2_lag)

    out += c[i1+1, i2+1]*evalBspline3near(x1_adv)*evalBspline3near(x2_adv)
    out += c[i1+2, i2+1]*evalBspline3far(x1_adv)*evalBspline3near(x2_adv)
    out += c[i1+1, i2+2]*evalBspline3near(x1_adv)*evalBspline3far(x2_adv)
    out += c[i1+2, i2+2]*evalBspline3far(x1_adv)*evalBspline3far(x2_adv)

    return out
end

function evalBspline3far(x::T) where T <: Real

    return 1/6 * (1-x)^3 # note this isn't (2-abs(x)), since we're using xk_adv.
    # taken from https://github.com/JuliaMath/Interpolations.jl/blob/master/src/b-splines/cubic.jl
end

function evalBspline3near(x::T) where T <: Real
    abs_x = abs(x)

    return 2/3 - x^2 + 0.5*x^3
end

# with respect to the second coordinate.
function cubicitp2Dexplicit2(c, x::Vector{T}, N::Vector{Int}) where T <: Real

    @assert one(T) <= x[1] <= convert(T,N[1])
    @assert one(T) <= x[2] <= convert(T,N[2])

    # determine anchor, 𝑖 = (i1,i2)
    δx1 = abs(x[1] - floor(x[1]))
    i1 = round(Int, x[1]-δx1)

    δx2 = abs(x[2] - floor(x[2]))
    i2 = round(Int, x[2]-δx2)

    x1_lag = x[1]-i1
    x2_lag = x[2]-i2

    x1_adv = 1 - x1_lag
    x2_adv = 1 - x2_lag

    a3 = zero(T)
    a2 = zero(T)
    a1 = zero(T)
    a0 = zero(T)

    k = i2

    ### x2_lag.

    # c[i1-1, i2-1]*evalBspline3far(x1_lag)*evalBspline3far(x2_lag)
    a3, a2, a1, a0 = getmonomialcoefficientsBspline3farlag(a3, a2, a1, a0,
                        c[i1-1, i2-1]*evalBspline3far(x1_lag), k)

    # c[i1, i2-1]*evalBspline3near(x1_lag)*evalBspline3far(x2_lag)
    a3, a2, a1, a0 = getmonomialcoefficientsBspline3farlag(a3, a2, a1, a0,
                        c[i1, i2-1]*evalBspline3near(x1_lag), k)

    # c[i1-1, i2]*evalBspline3far(x1_lag)*evalBspline3near(x2_lag)
    a3, a2, a1, a0 = getmonomialcoefficientsBspline3nearlag(a3, a2, a1, a0,
                        c[i1-1, i2]*evalBspline3far(x1_lag), k)

    # c[i1, i2]*evalBspline3near(x1_lag)*evalBspline3near(x2_lag)
    a3, a2, a1, a0 = getmonomialcoefficientsBspline3nearlag(a3, a2, a1, a0,
                        c[i1, i2]*evalBspline3near(x1_lag), k)

    ##
    # c[i1-1, i2+1]*evalBspline3far(x1_lag)*evalBspline3near(x2_adv)
    a3, a2, a1, a0 = getmonomialcoefficientsBspline3nearadv(a3, a2, a1, a0,
                        c[i1-1, i2+1]*evalBspline3far(x1_lag), k)

    # c[i1, i2+1]*evalBspline3near(x1_lag)*evalBspline3near(x2_adv)
    a3, a2, a1, a0 = getmonomialcoefficientsBspline3nearadv(a3, a2, a1, a0,
                        c[i1, i2+1]*evalBspline3near(x1_lag), k)

    # c[i1-1, i2+2]*evalBspline3far(x1_lag)*evalBspline3far(x2_adv)
    a3, a2, a1, a0 = getmonomialcoefficientsBspline3faradv(a3, a2, a1, a0,
                        c[i1-1, i2+2]*evalBspline3far(x1_lag), k)

    # c[i1, i2+2]*evalBspline3near(x1_lag)*evalBspline3far(x2_adv)
    a3, a2, a1, a0 = getmonomialcoefficientsBspline3faradv(a3, a2, a1, a0,
                        c[i1, i2+2]*evalBspline3near(x1_lag), k)

    ##
    # c[i1+1, i2-1]*evalBspline3near(x1_adv)*evalBspline3far(x2_lag)
    a3, a2, a1, a0 = getmonomialcoefficientsBspline3farlag(a3, a2, a1, a0,
                        c[i1+1, i2-1]*evalBspline3near(x1_adv), k)

    # c[i1+2, i2-1]*evalBspline3far(x1_adv)*evalBspline3far(x2_lag)
    a3, a2, a1, a0 = getmonomialcoefficientsBspline3farlag(a3, a2, a1, a0,
                        c[i1+2, i2-1]*evalBspline3far(x1_adv), k)

    # c[i1+1, i2]*evalBspline3near(x1_adv)*evalBspline3near(x2_lag)
    a3, a2, a1, a0 = getmonomialcoefficientsBspline3nearlag(a3, a2, a1, a0,
                        c[i1+1, i2]*evalBspline3near(x1_adv), k)

    # c[i1+2, i2]*evalBspline3far(x1_adv)*evalBspline3near(x2_lag)
    a3, a2, a1, a0 = getmonomialcoefficientsBspline3nearlag(a3, a2, a1, a0,
                        c[i1+2, i2]*evalBspline3far(x1_adv), k)

    ##
    # c[i1+1, i2+1]*evalBspline3near(x1_adv)*evalBspline3near(x2_adv)
    a3, a2, a1, a0 = getmonomialcoefficientsBspline3nearadv(a3, a2, a1, a0,
                        c[i1+1, i2+1]*evalBspline3near(x1_adv), k)

    # c[i1+2, i2+1]*evalBspline3far(x1_adv)*evalBspline3near(x2_adv)
    a3, a2, a1, a0 = getmonomialcoefficientsBspline3nearadv(a3, a2, a1, a0,
                        c[i1+2, i2+1]*evalBspline3far(x1_adv), k)

    # c[i1+1, i2+2]*evalBspline3near(x1_adv)*evalBspline3far(x2_adv)
    a3, a2, a1, a0 = getmonomialcoefficientsBspline3faradv(a3, a2, a1, a0,
                        c[i1+1, i2+2]*evalBspline3near(x1_adv), k)

    # c[i1+2, i2+2]*evalBspline3far(x1_adv)*evalBspline3far(x2_adv)
    a3, a2, a1, a0 = getmonomialcoefficientsBspline3faradv(a3, a2, a1, a0,
                        c[i1+2, i2+2]*evalBspline3far(x1_adv), k)

    return a3, a2, a1, a0
end

# monomial coefficients for evalBspline3far(adv)
## for β_far(x_k,adv): expand ((x-k)^3)/6
# x^3/6 -k*x^2/2 +k^2*x/2 -k^3/6
function getmonomialcoefficientsBspline3faradv(a3, a2, a1, a0,
                                                B::T,
                                                k::Int)::Tuple{T,T,T,T} where T <: Real
    b3 = 1/6
    b2 = -k/2
    b1 = k^2/2
    b0 = -k^3/6

    return B*b3+a3, B*b2+a2, B*b1+a1, B*b0+a0
end

# monomial coefficients for evalBspline3far(lag)
## for β_far(x_k,lag): expand ((1-(x-k))^3)/6
# k^3/6 - k^2*x/2 + k^2/2 + k*x^2/2 - k*x + k/2 - x^3/6 + x^2/2 - x/2 + 1/6
function getmonomialcoefficientsBspline3farlag(a3, a2, a1, a0,
                                                B::T,
                                                k::Int)::Tuple{T,T,T,T} where T <: Real
    b3 = -1/6
    b2 = (1 + k)/2
    b1 = -(1 + 2*k + k^2)/2
    b0 = (1 + 3*k + 3*k^2 + k^3)/6

    return B*b3+a3, B*b2+a2, B*b1+a1, B*b0+a0
end

# monomial coefficients for evalBspline3near(adv)
## for β_near(x_k,adv): expand 2/3 -(1-(x-k))^2+((1-(x-k))^3)/2
# 2/3 + k^3/2 - 3*k^2*x/2 + k^2/2 + 3*k*x^2/2 -k*x -k/2 -x^3/2 +x^2/2 +x/2 -0.5
function getmonomialcoefficientsBspline3nearadv(a3, a2, a1, a0,
                                                B::T,
                                                k::Int)::Tuple{T,T,T,T} where T <: Real
    b3 = -0.5
    b2 = (1 + 3*k )/2
    b1 = (1 - 2*k - 3*k^2)/2
    b0 = (-1 -k + k^2 + k^3)/2 + 2/3

    return B*b3+a3, B*b2+a2, B*b1+a1, B*b0+a0
end

# monomial coefficients for evalBspline3near(lag)
## for β_near(x_k,lag): expand 2/3 -(x-k)^2 +((x-k)^3)/2
# -k^3/2 + 3*k^2*x/2 -k^2 -3*k*x^2/2 +2*k*x +x^3/2 -x^2 +2/3
function getmonomialcoefficientsBspline3nearlag(a3, a2, a1, a0,
                                                B::T,
                                                k::Int)::Tuple{T,T,T,T} where T <: Real
    b3 = 0.5
    b2 = -1 - 1.5*k
    b1 = 2*k + 1.5*k^2
    b0 = 2/3 - k^2 - k^3/2

    return B*b3+a3, B*b2+a2, B*b1+a1, B*b0+a0
end

### drawing RVs.
function drawcategorical(w::Vector{T})::Int where T <: Real
    x = rand()

    tmp = 0.0
    for k = 1:length(w)-1
        if x < w[k]+tmp
            return k
        end
        tmp += w[k]
    end

    return length(w)
end

### iterated Brownian Bridge

function Bernoullipolynomial0(x::T)::T where T <: Real
    return one(T)
end

function Bernoullipolynomial1(x::T)::T where T <: Real
    return x-0.5
end

function Bernoullipolynomial2(x::T)::T where T <: Real
    return x^2 -x +1/6
end

function Bernoullipolynomial3(x::T)::T where T <: Real
    return x^3 -1.5*x^2 + 0.5*x
end

function Bernoullipolynomial4(x::T)::T where T <: Real
    return x^4 -2*x^3 +x^2 -1/30
end

function Bernoullipolynomial5(x::T)::T where T <: Real
    return x^5 -2.5*x^4 +5/3*x^3 -1/6*x
end

function Bernoullipolynomial6(x::T)::T where T <: Real
    return x^6 -3*x^5 +5/2*x^4 -0.5*x^2 +1/42
end

### visualization helpers.

# options for warp map samples.
# this is the case when the target density is unknown, but realizations are available.
function getwarpmapsamplecustom( y::Array{T,D},
                            ω_set,
                            pass_band_factor) where {T,D}
    #
    N_bands = length(ω_set)

    Y = y

    #### Split-band analysis.
    ϕY, ψY = SignalTools.runsplitbandanalysis(Y, ω_set, SignalTools.getGaussianfilters)
    ηY = SignalTools.runbandpassanalysis(Y, ω_set, pass_band_factor, SignalTools.getGaussianfilters)

    # #### Riesz transform on the different filtered signals.
    # H, ordering = gethigherorderRTfilters(Y,order)
    #
    # 𝓡ϕY = collect( RieszAnalysisLimited(ϕY[s],H) for s = 1:N_bands)
    # 𝓡ψY = collect( RieszAnalysisLimited(ψY[s],H) for s = 1:N_bands)
    # 𝓡ηY = collect( RieszAnalysisLimited(ηY[s],H) for s = 1:N_bands)

    ϕ_set = ηY

    ϕ = reduce(+,ϕ_set)./N_bands

    return ϕY, ψY, ηY
end

function getwarpmaplinear(ϕ::Array{T,D}) where {T,D}

    itp_ϕ = Interpolations.interpolate(ϕ,
                Interpolations.BSpline(Interpolations.Linear()))

    etp_ϕ = Interpolations.extrapolate(itp_ϕ, Interpolations.Line())

    return etp_ϕ
end

# 1-D CDF: integrate from -Inf to b.
# f is the probability density function.
# check the Jacobian: derivative of b-x/(1-x)-0.5.
# check the b bound becomes -1: evaluate b-x/(1-x)-0.5, x = -1
# for use with HCubature, so vector-valued input.
function get1DCDFintegrand(b::T, f::Function) where T <: Real

    return tt->f(b - tt[1]/(1-tt[1]) -0.5)*(1/(1-tt[1])^2)
end

function getKRCDFintegrand(b::T, f::Function) where T <: Real

    return tt->f(b - tt[end]/(1-tt[end]) -0.5)*(1/(1-tt[end])^2) # Sign flip because we have the bounds reversed.
end

# compute the normalizing constant.
# for use with HCubature, so vector-valued input.
function get1DZintegrand(f::Function) where T <: Real

    return tt->f(tt[1]/(1-tt[1]^2))*((1+tt[1]^2)/(1-tt[1]^2)^2)
end

# computes 1D normalizing constant in the last input component of f.
function getKRZintegrand(f::Function) where T <: Real

    return tt->f(tt[end]/(1+tt[end])^2)*(1+tt[end]^2)/(1-tt[end]^2)^2
end

# computes 1D normalizing constant in the last input component of f.
function getKRZintegrandlogspace(f::Function) where T <: Real

    return tt->exp( log(f(tt[end]/(1+tt[end])^2)) + log(1+tt[end]^2) - 2*log(1-tt[end]^2) )
end

# returns positions, weights. Not working.
function gettanhsinhquad(N::Int, m::Int)::Tuple{Vector{Float64},Vector{Float64}}
    half_π = π/2

    j_range = -N:N
    h = 1/2^m

    g = tt->tanh(half_π*sinh(tt))
    dg = tt->half_π *cosh(tt)*sech(half_π*sinh(tt))^2 # derivative of tanh(pi/2*sinh(t))

    x = collect( g(h*j) for j in j_range )
    w = collect( dg(h*j) for j in j_range )

    return x, w
end

# evaluates a tanh-sinh quadature integral.
function computetsq(x, w, f, m)
    N_nodes = length(w)
    h = 1/2^m

    return h*sum( w[i]*f(x[i]) for i = 1:N_nodes )
    #return sum( w[i]*f(x[i]) for i = 1:N_nodes )
end

# coordinates.
# x ∈ prod( [a[d], b[d]] for d = 1:D ).
# same sampling interval for each dimension.
# M[d] is the number of samples in dimension d.
# out ∈ prod( 1:M[d] for d = 1:D )
function convert2itpindex(  x::Vector{T},
                            a::Vector{T},
                            b::Vector{T},
                            M::Vector{Int})::Vector{T} where T <: Real

    D = length(x)
    @assert D == length(a) == length(b)

    out = Vector{T}(undef,D)
    for d = 1:D
        offset = b[d]-a[d]
        u = (x[d]-a[d])/offset # ∈ [0,1]
        i = u*(M[d]-1) # ∈ [0,M-1]
        out[d] = i + 1
    end

    return out
end
# x =[-0.01; 2]
# a = [-0.1; -1 ]
# b = [0.0; 3]
# M = [10; 5]
# p = convert2itpindex([-0.01; 3], a, b, M) # should be [9;5]

# marginalize over ℝ for the dimensions specified in u_dims.
# Uses the change-of-variable here:
# https://github.com/stevengj/cubature/blob/master/README.md#infinite-intervals
# function marginalize(f::Function, u_dims::Vector{Float64}, ϵ = 1e-4)
#
# (1+tt[d]^2)/(1-tt[d]^2)^2)
#
#     d = u[1]
#     integrand_func = tt->f(tt[d]/(1-tt[d]^2)
#     for d in u[1:end]
#         integrand_func = tt->integrand_func(tt[d]/(1-tt[d]^2)
#     end
#
#     x_min = -one(Float64)
#     x_max = one(Float64)
#     (val,err) = HCubature.hcubature(integrand_func, [x_min], [x_max];
#                     norm = LinearAlgebra.norm, rtol = sqrt(eps(Float64)), atol = 0,
#                     maxevals = typemax(Int), initdiv=1)
#
#     return val, err
# end

function iterativesin()
    N = 3
    f_array = Vector{Function}(undef,N)

    f_array[1] = xx->sin(xx)
    for i = 2:N
        f_array[i] = xx->f_array[i-1](sin(xx))
    end

    return f_array[end]
end

# # f is the joint pdf.
# # marginalizes over variables N:D, where D is the total number of variables.
# # assumes the N variables have support ℝ.
# function getintegrandformarginalization(v::Vector{T},
#                                         f::Function,
#                                         N::Int,
#                                         D::Int) where T <: Real
#     @assert 0 < N < D # at least one variable remains.
#     @assert length(v) == D-N
#
#     g = tt->tt/(1-tt^2)
#     u = xx->collect( g(xx[d]) for d = N:D )
#
#     J = tt->( (1+tt^2)/(1-tt^2)^2 )
#
#     integrand_func = xx->f( [v; u(xx)] )*prod( J(xx[d]) for d = N:D )
#
#     return integrand_func
# end
#
# # f is joint pdf.
# # based on numerical integration.
# function evalmarginalpdf(f::Function, v::Vector{Float64}, N::Int, D::Int, max_integral_evals::Int = 1000)
#
#     integrand_func = getintegrandformarginalization(v, f, N, D)
#
#     x_min = collect( -1.0 for n = N:D )
#     x_max = collect( 1.0 for n = N:D )
#     (val,err) = HCubature.hcubature(integrand_func, x_min, x_max;
#                     norm = LinearAlgebra.norm, rtol = sqrt(eps(Float64)), atol = 0,
#                     maxevals = max_integral_evals, initdiv = 1)
#     #
#     return val
# end

# f is the joint pdf.
# marginalizes over variables 1:N. D is the total number of variables.
# assumes the N variables have support ℝ.
function getintegrandformarginalization(v::Vector{T},
                                        f::Function,
                                        N::Int,
                                        D::Int) where T <: Real
    @assert N < D # at least one variable remains.
    @assert length(v) == D-N

    g = tt->tt/(1-tt^2)
    u = xx->collect( g(xx[d]) for d = 1:N )

    J = tt->( (1+tt^2)/(1-tt^2)^2 )

    integrand_func = xx->f( [u(xx); v] )*prod( J(xx[d]) for d = 1:N )
    #integrand_func = xx->f( [u(xx); u(v)] )*prod( J(xx[d]) for d = 1:length(xx) )*prod( J(v[d]) for d = 1:length(v) )

    return integrand_func
end
# f is joint pdf.
# based on numerical integration.
function evalmarginalpdf(f::Function, v::Vector{Float64}, N::Int, D::Int, max_integral_evals::Int = 1000)

    integrand_func = getintegrandformarginalization(v, f, N, D)

    x_min = collect( -1.0 for n = 1:N )
    x_max = collect( 1.0 for n = 1:N )
    (val,err) = HCubature.hcubature(integrand_func, x_min, x_max;
                    norm = LinearAlgebra.norm, rtol = sqrt(eps(Float64)), atol = 0,
                    maxevals = max_integral_evals, initdiv = 1)
    #
    return val
end

# f is the joint pdf.
# marginalizes over the last N variables, where D is the total number of variables.
# assumes the N variables have support ℝ.
function getintegrandformarginalizationrev(v::Vector{T},
                                        f::Function,
                                        N::Int,
                                        D::Int) where T <: Real
    @assert 0 < N < D
    @assert length(v) == D- length((D-N+1):D)

    g = tt->tt/(1-tt^2)
    u = xx->collect( g(xx[d]) for d = 1:length(xx) )

    J = tt->( (1+tt^2)/(1-tt^2)^2 )

    integrand_func = xx->f( [v; u(xx)] )*prod( J(xx[d]) for d = 1:length(xx) )

    return integrand_func
end
# f is joint pdf.
# based on numerical integration.
function evalmarginalpdfrev(f::Function, v::Vector{Float64}, N::Int, D::Int, max_integral_evals::Int = 1000)

    integrand_func = getintegrandformarginalizationrev(v, f, N, D)

    x_min = collect( -1.0 for n = (D-N+1):D )
    x_max = collect( 1.0 for n = (D-N+1):D )
    (val,err) = HCubature.hcubature(integrand_func, x_min, x_max;
                    norm = LinearAlgebra.norm, rtol = sqrt(eps(Float64)), atol = 0,
                    maxevals = max_integral_evals, initdiv = 1)
    #
    return val
end

# f is the joint pdf.
# marginalizes over the first N variables.
# assumes the N variables have support ℝ.
function getintegrandforZ(  f::Function,
                            D::Int) where T <: Real

    g = tt->tt/(1-tt^2)
    u = xx->collect( g(xx[d]) for d = 1:D )

    J = tt->( (1+tt^2)/(1-tt^2)^2 )

    integrand_func = xx->f( u(xx) )*prod( J(xx[d]) for d = 1:D )

    return integrand_func
end

# f is joint pdf.
# based on numerical integration.
function evalZ9(f::Function, D::Int, max_integral_evals::Int = 10000)

    integrand_func = getintegrandforZ(f, D)

    x_min = collect( -1.0 for d = 1:D )
    x_max = collect( 1.0 for d = 1:D )
    (val,err) = HCubature.hcubature(integrand_func, x_min, x_max;
                    norm = LinearAlgebra.norm, rtol = sqrt(eps(Float64)), atol = 0,
                    maxevals = max_integral_evals, initdiv = 1)
    #
    return val, err
end

# Uses the change-of-variable here:
# https://github.com/stevengj/cubature/blob/master/README.md#infinite-intervals
function evalindefiniteintegral(f::Function, max_integral_evals = typemax(Int))

    integrand_func = tt->(f(tt[1]/(1-tt[1]^2))*(1+tt[1]^2)/(1-tt[1]^2)^2)
    x_min = -one(Float64)
    x_max = one(Float64)
    (val,err) = HCubature.hcubature(integrand_func, [x_min], [x_max];
                    norm = LinearAlgebra.norm, rtol = sqrt(eps(Float64)), atol = 0,
                    maxevals = max_integral_evals, initdiv=1)

    return val, err
end

# for p(A | B) as a function of f(A,B).
function setupCDFkernel(hyperparam_A::T,
                        hyperparam_canonical_B::T,
                        ϕ_map::Function)::CPDF1DKernelType{T,Spline34KernelType{T}} where T <: Real

    θ_A = RationalQuadraticKernelType(hyperparam_A)

    θ_canonical = Spline34KernelType(hyperparam_canonical_B) # for adaptive kernel.
    θ_B = AdaptiveKernelType(θ_canonical, xx->ϕ_map(xx))

    return CPDF1DKernelType(θ_A, θ_B)
end


using Base.Test

# test_meshwarp()
img = reshape(float(collect(1:121).%2), 11, 11) # 11x11 checkerboard
triangles = [1 2 3]

tform = [1 0 0;
        0 1 0;
        0 0 1]
src = [0.0 0.0;
        0.0 10.0;
        10.0 0.0]
dst = warp_pts(tform, src)
offset = [0,0]
img_warped, warped_offset = meshwarp(img, src, dst, triangles, offset)
# @test_approx_eq_eps img_warped img 1e-5
@test warped_offset == [0,0]

tform = [1 0 0;
        0 1 0;
        0 0 1]
src = [0.0 0.0;
        0.0 10.0;
        10.0 0.0]
dst = warp_pts(tform, src)
offset = [5,10]
img_warped, warped_offset = meshwarp(img, src, dst, triangles, offset)
@test_approx_eq_eps img_warped zeros(11,11) 1e-10
@test warped_offset == [0,0]

tform = [1 0 0;
        0 1 0;
        0 0 1]
src = [0.0 0.0;
        0.0 20.0;
        20.0 0.0]
dst = warp_pts(tform, src)
offset = [5,10]
img_warped, warped_offset = meshwarp(img, src, dst, triangles, offset)
@test img_warped[6,11] == 1.0
@test warped_offset == [0,0]

src = [2.0 2.0;
        10.0 2.0;
        6.0 10.0]
dst = [4.0 2.0;
        8.0 2.0;
        6.0 10.0]
offset = [0, 0]
img_warped, warped_offset = meshwarp(img, src, dst, triangles, offset)
@test warped_offset == [4,2]

# test_poly2mask()
# vertices=[1.5 2.5; 2.5 1.5; 3.5 4.5]
# img, off = poly2mask(vertices[:,1],vertices[:,2])
# """ off should be [1,1] and img should be:
# false  false  false  false  false
# false   true   true  false  false
# false  false   true   true  false
# false  false  false  false  false
# """
# vertices=[1.5 2.4; 2.5 1.5; 3.5 4.5]
# img, off = poly2mask(vertices[:,1],vertices[:,2])
# """ off should be [1,1] and img should be:
# false  false  false  false  false
# false   true  false  false  false
# false  false   true  false  false
# false  false  false  false  false,
# """
# vertices=[1 2; 2.5 1.5; 3.5 4.5]
# img, off = poly2mask(vertices[:,1],vertices[:,2])
# """ corner case: off should be [1,1] and img should be:
# false   true  false  false  false
# false   true   true  false  false
# false  false   true   true  false
# false  false  false  false  false
# """
# p=plot(x=vertices[:,1],y=vertices[:,2])
# display(p)
# img,off


using Base: Bool
using LinearAlgebra
using SparseArrays

# functions

function unity(i,n)
    E = spzeros(Bool,n,n)
    E[i] = true
    return E
end

# xs = [rand(3,4) for _ in 1:5]
# reduce((x,y) -> cat(x,y,dims=3), A)
# idea: https://github.com/JuliaLang/julia/pull/37196
# reduce(cat(dims=4), [ones(2,3) for _ in 1:5]) |> size
#`reduce(vcat, [rand(2,3) for _ = 1:4]): A[i][a,b] -> A[(i-1)*len_a + a, b]`
#which is the same as restructuring to` B[i,a][b] via B = [A[i][a,:] for i = 1:dim_i, a = 1:dim_a] `
#and then taking `reduce(hcat,B'[:]')'`
#`reduce(hcat, [rand(2,3) for _ = 1:4]): A[i][a,b] -> A[a, (i-1)*len_b + b]`

function partial_trace(dim_1, dim_2, traced_system)
# traces out the system indicated by traced_system
# for C = kron(A,B), tracing over the first system (traced_system = 1);
# yields: reshape(p1*C[:], 4,4)/tr(A)  == B
#
# in colex notation:
# traced_system == 1 represents a trace over index 2 and 4 and
# traced_system == 2 represents a trace over index 1 and 3 and
    if traced_system == 1
        T = [kron(sparse(I, dim_1,dim_1), unity(i,dim_2))[:] for i = 1:dim_2^2]
        col = reduce(vcat,[findnz(T[i])[1] for i = 1:dim_2^2])
        row = kron(1:dim_2^2, ones(Int64,dim_1))
        return sparse(row, col, trues(dim_2^2*dim_1))
    else
        T = [kron(unity(i, dim_1), sparse(I, dim_2, dim_2))[:] for i = 1:dim_1^2]
        col = reduce(vcat, [findnz(T[i])[1] for i = 1:dim_1^2])
        row = kron(1:dim_1^2, ones(Int64,dim_2))
        return sparse(row, col, trues(dim_1^2*dim_2))
        #return    Tr_B = [reshape(kron(reshape(unity(i,dim_1), dim_1, dim_1), I(dim_2)), 1, dim_2^2 * dim_1^2)[1,j] for i = 1:dim_1^2, j =1:dim_1^2*dim_2^2]
    end
end

function choi(A, dim_1)
mm, nn = size(A)
n = mm ÷ dim_1
return reshape(permutedims(reshape(A,n,dim_1,n,dim_1),(1,3,2,4)), mm,mm)
end

function choi_rect(A,dim_1, dim_2, dim_3, dim_4)
return reshape(permutedims(reshape(A, dim_1, dim_2, dim_3, dim_4), (1,3,2,4)), dim_1*dim_3, dim_2*dim_4)
end

function ptr(ρ, dim_1, traced_system = 1)
    # get partial trace superoperator
    N,M = size(ρ)
    dim_2 = N÷dim_1
    PT = partial_trace(dim_1, dim_2, traced_system)
    new_dim = traced_system == 1 ? dim_2 : dim_1
    # apply superoperator to vectorized density matrix and reshape to new dimensions
    σ = reshape(PT * ρ[:], new_dim, new_dim)
    return σ
end

### apply split
function factorize(A, dim_1, dim_2)
    σ = prt(A, dim_1, dim_2, 2)
    ρ = prt(A, dim_1, dim_2, 1)
    A_tensor_prod = kron(σ, ρ)
    rho_c = A - A_tensor_prod
    return A_tensor_prod, rho_c
end

### density matrix

function density_matrix(n, neg = 0, real_valued = false)
    #trace 1 remains preserved!
    dia = rand(n)
    if neg > 0
        dia[rand(1:n, neg)] .*= -1
    end
    dia = dia/sum(dia)
    if real_valued
        temp = rand(n,n)
    else
        temp = rand(n,n) + rand(n,n) * im
    end
    ew, ev = eigen(Hermitian(temp'*temp))
    d_mat = Hermitian(ev* Diagonal(dia) * ev')
    return d_mat
end

function choi2kraus(C)
    dim_1,dim_2 = size(C)
    N= Int64(sqrt(dim_1))
    K = [zeros(ComplexF64, N,N) for _=1:N^2]
    ew,ev = eigen(C)
    for k = 1:N^2
        K[k] = reshape(conj.(ev[:,k]),N,N) .* sqrt(abs(ew[k]))
        if ew[k] < 0
            print("Warning, negative eigenvalue, k = " * string(k))
        end
    end
    return K
end

function kraus2choi(K)
    dim1,dim2 = size(K[1])
    C = zeros(ComplexF64, dim1^2,dim1^2)
    for k = 1:length(K)
        C += conj.(K[k])[:] * transpose(K[k][:])
    end
    return C
end
### ======================== functions end =======================

dim_1 = 2
dim_2 = 3
m = 7
n = 3

A = rand(m,m)
B = rand(n,n)

C = kron(A,B)

A_sup = kron(A[:], I(n^2))
C_sup = choi(C,m)

sum(abs.(A_sup * B[:] - C_sup[:]))

A_sup * B[:] - C[:]
reshape(A_sup * B[:],n*m,n*m) - C_sup
sum(abs.(ans[:]))

# create superoperator without applying Choi on the right hand side

G = reshape(permutedims(reshape(kron(A[:], I(n^2)), (n,n,m,m,n^2)), (1,3,2,4,5)), n^2*m^2,n^2)

reshape(G*B[:], m*n, m*n)
C
sum(abs.(G * B[:] - C[:]))

kron_fac_B = reshape(permutedims(reshape(kron(I(m^2), B[:]), (n,n,m,m,m^2)), (1,3,2,4,5)), n^2*m^2, m^2)
norm(reshape(kron_fac_B * A[:], n*m, n*m) - C)

function I_rect(n1,n2)
    if n1 >= n2
        temp = I(n1)
        temp = temp[:,1:n2]
    else
        temp = I(n2)
        temp = temp[1:n1,:]
    end
    return temp
end

m1 = 3
m2 = 4

n1 = 5
n2 = 6

A = rand(m1,m2)
B = rand(n1,n2)

C = kron(A,B)

#G = reshape(permutedims(reshape(kron(A[:], I(n1*n2)), (n1,n2,m1,m2,n1*n2)), (1,3,2,4,5)), n1*n2*m1*m2,n1*n2)

S_left = reshape(permutedims(reshape(kron(I(n1*n2), A[:]), (m1,m2,n1,n2,n1*n2)), (3,1,4,2,5)), prod([m1,m2,n1,n2]), n1*n2)
norm(S_left * B[:] - C[:])

S_right = reshape(permutedims(reshape(kron(B[:], I(m1*m2)), (m1,m2,n1,n2,m1*m2)), (3,1,4,2,5)), prod([m1,m2,n1,n2]), m1*m2)
norm(S_right * A[:] - C[:])

norm(kron(A[:], I(n1*n2)) * B[:] - choi_rect(C, n1,m1, n2, m2)[:])

norm(kron(I(m1*m2), B[:]) * A[:] - choi_rect(C,n1,m1,n2,m2)[:])

norm(kron(A'[:], I(n1*n2)) * B'[:] - C'[:])

# elegant version to show right hand side as a Choi matrix for first and second factor:
norm(kron(I(n1*n2) ,A[:]) * B[:] - choi_rect(C, n1, m1, n2, m2)'[:])
norm(kron(B[:], I(m1*m2)) * A[:] - choi_rect(C, n1, m1, n2, m2)'[:])

# check embeddings

A = rand(2)
B = rand(3)
C = rand(4)
D = rand(5)

Q = kron(A, B')
kron(kron(A, B'), ones(4,5))

kron(kron(A,B'), kron(C,D')) - kron(kron(A, B'), I(4)) * kron(I(4), kron(C,D'))

Q = kron(A, C')
kron(kron(A, ones(3)'), kron(C, ones(5)'))- kron(Q'[:], ones(3*5)')

Q = kron(A, D')
kron(kron(A, ones(3)'), kron(ones(4), D'))-kron(kron(ones(3)', Q), ones(4))

Q = kron(B, C')
kron(kron(ones(2), B'), kron(C, ones(5)')) - kron(kron(ones(2),Q'), ones(5)')

Q = kron(B, D')
kron(kron(ones(2), B'), kron(ones(4), D'))-kron(ones(2*4), Q[:]')

# multiplication table for nxn unity base matrices

n = 4
A = zeros(Int64,n^2,n^2)
function unity(i,n)
    E = spzeros(Bool,n,n)
    E[i] = true
    return E
end

for i = 1:n^2, j=1:n^2
    L = findall((unity(i,n)*unity(j,n))[:] .> 0)
    if isempty(L)
        A[i,j] = 0
    else
        A[i,j] = L[1]
    end
end
A

# try to check \sum_{i != j} a_ii a_jj - abs(a_ij)^2 == sum_{i \neq j} \lambda_i lambda_j

A = density_matrix(10,0);
#Diagonal(real(A))
ew,ev = eigen(A);
L1 = [ A[i,i]*A[j,j] - abs(A[i,j])^2 for i = 1:5, j = 1:5];
L2 = [ew[i] * ew[j] * (i != j) for i = 1:5, j=1:5];

sum(real(L1) .< 0)
real(ew)'

# can every density matrix be represented as the sum of tensor products of valid density matrices?

A = density_matrix(4)
B = density_matrix(3)

ew, ev = eigen(kron(A,B))
sum(ew .< 0)

a = rand(4)

# the Kraus decomposition for initially correlated dynamical maps:
A6 = density_matrix(6)
A2 = density_matrix(2)
A3 = density_matrix(3)
A4 = density_matrix(4)

A = density_matrix(9)

n_runs = 80

tr_weight = zeros(n_runs)

A =  kron(A3, A4)
for k = 1:n_runs
    T1, C1 = factorize(A./tr(A), 2,6)

    A, M1 = split(C1)
    tr_weight[k] = tr(A)
end
tr_weight
sum(cumprod(tr_weight))

# entangled state

A = [1 0 0 1; 0 0 0 0; 0 0 0 0; 1 0 0 1]/2
t,c = factorize(A/tr(A),2,2)
p,m = split(c)

ω = tr(p)
p = p/tr(p)
m = m/tr(m)
round)
function split(rho)
    ew, ev = eigen(rho)
    rho_p = Hermitian(ev[:,ew .> 0] * Diagonal(ew[ew.> 0]) * ev[:,ew .> 0]')
    rho_m = Hermitian(ev[:,ew .< 0] * Diagonal(-ew[ew.< 0]) * ev[:,ew .< 0]')
    return rho_p, rho_m
end

P1 = partial_trace(2,3,1)
P2 = partial_trace(2,3,2)

A1 = reshape(P2*A[:],2,2)
A2 = reshape(P1*A[:],3,3)

rho_c = A - kron(A1,A2)
rho_c = Hermitian(rho_c)

ew,ev = eigen(rho_c)

rho_p = Hermitian(ev[:,ew .> 0] * Diagonal(ew[ew.> 0]) * ev[:,ew .> 0]')
rho_m = Hermitian(ev[:,ew .< 0] * Diagonal(-ew[ew.< 0]) * ev[:,ew .< 0]')

# now decompose both again, maybe it does converge

A = rand(4,4)
B = rand(4,4)

kron(A,B)' - kron(A', B')

# Choi to Kraus representation

A = density_matrix(16) * 4

ew, ev = eigen(A)
K = [zeros(ComplexF64,4,4) for _ in 1:16]
for k = 1:16
    K[k] = reshape(ev[:,k], 4,4) .* sqrt(ew[k])
end

Test = zeros(ComplexF64,4,4)
for k=1:16
    Test = Test + K[k] * K[k]'
end
Test
tr(Test)
B = density_matrix(4)
B2 = zeros(ComplexF64,4,4)
for k = 1:16
    B2 += K[k] * B * K[k]'
end
tr(B)
tr(B2)

A = rand(4,4)
A = A

exp(-im*A) * exp(-im*A)'

# general way to construct a valid quantum operation:
# find A \in C^{d^2 x d^2}
# with A^\dag = A and A>0
# then A2 = A + (I - P_13 A) \kron I/d is a valid Choi matrix

# create choi matrix directly fulfilling the trace condition

A = density_matrix(16)*4
P13 = partial_trace(4,4,2) #tracing out the second system - flip to switch to colexicographic
P24 = partial_trace(4,4,1)

A2 = A + kron(I(4) - reshape(P13*A[:], 4,4), I(4)/4)

ew, ev = eigen(A2); ew
round(sum(abs.(ew)) - tr(A2), digits=14)

reshape(P13*A2[:], 4,4)
# unital map: Tr_24 yields I(N)
A3 = A2 + kron(I(4)/4, I(4) - reshape(P24*A[:],4,4))

round.(reshape(P13*A3[:], 4,4), digits = 14)

round.(reshape(P24*A3[:], 4,4), digits = 14)

# get Kraus operators from A2 from eigenvectors
#(complex conjugation can be ignored since C is hermitian):
K = [zeros(ComplexF64, 4,4) for _=1:16]
for k = 1:16
    K[k] = reshape(ev[:,k],4,4) * sqrt(ew[k])

end
K = choi2kraus(A2)
id_test = zeros(ComplexF64,4,4)
for k = 1:16
    id_test += K[k]' * K[k]
end

id_test

C = kraus2choi(K)
sum(abs.(C - (A2)))
C
A2

#Ginibre matrix NxN
#G = randn(N,N)
# 1/(2 pi )^(N^2/2) * exp( -1/2 Tr(G * transpose(G)))

N= 1000
G = randn(N,N)
1/(2*pi)^(N^2/2) * exp(- 1/2 * tr(G * transpose(G)))

# create random Kraus operators from large unitary matrix
N = 4
A = rand(N^2*N, N^2*N)
ew, ev = eigen(Hermitian(A))
K = [zeros(ComplexF64, 4,4) for _=1:16]
for k=1:16
    K[k] = ev[1:4, (1:4) .+ (k-1)*4]'
end

# get the choi matrix from the Kraus operators
C = zeros(ComplexF64, 16,16)

for k = 1:16
    C = C + K[k]'[:] * K[k]'[:]'
end
ew, ev = eigen(C);
tr(C) - sum(ew)
sum(abs.(C - C'))

# get linear superoperator
D = choi(C,4)

p1 = partial_trace(4,4,1)

Test  = zeros(4,4)
reshape(sum(p1, dims = 1),4,4)

A = randn(12,16) + randn(12,16) * 0.1*im
C_ = A'*A/tr(A'*A)*4
w,v = eigen(C_); w

#C_ = density_matrix(16) * 4
C = (C_ + kron(I(4) - ptr(C_, 4, 2), I(4)/4))
tr(C)
ew, ev = eigen(C);ew

ptr(C,4,2)
reshape(P13 * C[:],4,4)

# Kraus operators of partial transpose

#transpose operator
T = [1 0 0 0; 0 0 1 0; 0 1 0 0; 0 0 0 1]

C = choi(T,2)
ew, ev = eigen(C)

K = choi2kraus(C)
A = density_matrix(2)
temp = zeros(ComplexF64,2,2)
for k = 1:4
    temp += K[k] * A * K[k]'
    if k == 1
        temp *=-1
    end
end
temp - A
A
Id_test  = zeros(ComplexF64,2,2)
for k = 1:4
 Id_test += K[k]' * K[k]
end
Id_test

K[1]' * K[1]

A = rand(4,4)

B = rand(3,3)

# tensor product as superoperator

# L(X) = A ⊗ X

A = rand(2,2)
X = rand(3,3)
# aim:
L * X[:] == kron(A,X)[:]

CL = kron(A[:], I(9))
L = reshape(permutedims(reshape(CL, 3,3,2,2,9),[1,3,2,4,5]), 36,9)
L * X[:]
kron(A,X)[:]

L2 = kron(kron(A,I(3))[:], I(3)[:]')
L

i = 16

n = [2,3,2,3]

a = [(i-1) % n[1]+1,
((i-1) % prod(n[1:2]) ÷ n[1]) + 1,
(i-1) % prod(n[1:3]) ÷ prod(n[1:2]) + 1,
(i-1) ÷ prod(n[1:3]) + 1]

a[1] + n[1] * (a[2]-1 + n[2] * (a[3] - 1 + n[3] * (a[4]-1)))

function choi_index(i, n)
    return  (i.-1) .% n[1] .+ 1 .+
    n[1] .* ((i.-1) .% prod(n[1:3]) .÷ prod(n[1:2]) .+
    n[2] .* ((i.-1) .% prod(n[1:2]) .÷ n[1] .+
    n[3] .* ((i.-1) .÷ prod(n[1:3]))))
end
T = reshape(1:36,2,3,2,3)


include("../src/mesh.jl")
include("../src/advection.jl")
include("../src/lagrange.jl")
# include("../src/lagrange2d.jl")
# include("../src/advection2d.jl")
include("../src/bspline.jl")
include("../src/bsplinelu.jl")
include("../src/bsplinefft.jl")
include("../src/interpolation.jl")

using LinearAlgebra
using Plots

function exact(tf::T, mesh1::UniformMesh{T}, mesh2::UniformMesh{T}) where {T}

    f = zeros(T, (mesh1.length, mesh2.length))
    for (i, x) in enumerate(mesh1.points), (j, y) in enumerate(mesh2.points)
        xn = cos(tf) * x + sin(tf) * y
        yn = -sin(tf) * x + cos(tf) * y
        f[i, j] = exp(-6 * ((xn)^2 + (yn + big"1.8")^2))
    end

    return f

end

function fctadv(interp, mesh1, mesh2, v1, v2, refdeb, refend, dt)
    adv_x1 = Advection1d(mesh1, interp)
    adv_x2 = Advection1d(mesh2, interp)

    f = copy(refdeb)
    ft = copy(refdeb)
    advection!(adv_x1, f, v1, tan(dt / 2))
    transpose!(ft, f)
    advection!(adv_x2, ft, v2, sin(dt))
    transpose!(f, ft)
    advection!(adv_x1, f, v1, tan(dt / 2))

    return norm(f - refend, Inf)
end

function fctmain(sz, dt::T, ordmax) where {T}

    mesh1 = UniformMesh(T(-5), T(5), sz)
    mesh2 = UniformMesh(T(-5), T(5), sz)

    refdeb = exact(zero(T), mesh1, mesh2)
    refend = exact(dt, mesh1, mesh2)

    v1 = -collect(mesh2.points)
    v2 = +collect(mesh1.points)

    indmax = div(ordmax - 3, 2) + 1

    y = ones(indmax, 3)
    x = zeros(Int64, indmax)

    labels = Array{String,2}(undef, 1, 3)
    labels[1, 1] = "Lagrange"
    labels[1, 2] = "B-Spline LU"
    labels[1, 3] = "B-Spline FFT"

    ind = 1

    for order = 3:2:ordmax
        println("order=$order")
        lag = Lagrange(BigFloat, order)
        splu = B_SplineLU(order, mesh1.length, zero(T))
        spfft = B_SplineFFT(order, mesh1.length, zero(T))

        y[ind, 1] = fctadv(lag, mesh1, mesh2, v1, v2, refdeb, refend, dt)
        y[ind, 2] = fctadv(splu, mesh1, mesh2, v1, v2, refdeb, refend, dt)
        y[ind, 3] = fctadv(spfft, mesh1, mesh2, v1, v2, refdeb, refend, dt)

        println("ind=$ind order=$order y[ind,:]=$(y[ind,:])")

        x[ind] = order

        ind += 1
    end

    p = Plots.plot(
        x,
        log10.(y),
        xlabel = "order",
        ylabel = "error",
        legend = :bottomleft,
        label = labels,
        marker = 2,
    )
    prec = precision(BigFloat)

    Plots.savefig(p, "out/result2_$(sz)_$(prec)_$(ordmax).pdf")
end

function fctinter(interp::InterpolationType, sz, dec)

    fct(x) = cos(2big(pi) * x)
    mesh = big.(1:sz) / sz
    deb = fct.(mesh)
    ref = fct.(mesh .+ dec / sz)
    res = zeros(BigFloat, sz)
    interpolate!(res, deb, dec, interp)
    return Float64(norm(ref - res, Inf))
end

function fctmain2(sz, ordmax)

    indmax = ordmax - 2

    y = ones(indmax, 2)
    x = zeros(Int64, indmax)

    labels = Array{String,2}(undef, 1, 2)
    labels[1, 1] = "Lagrange"
    labels[1, 2] = "B-Spline LU"
    #    labels[1,3] = "B-Spline FFT"

    ind = 1

    for order = 3:ordmax
        println("order=$order")
        lag = Lagrange(BigFloat, order)
        splu = B_SplineLU(order, sz, big"0.")
        #        spfft = B_SplineFFT(order, sz, big"0.")

        y[ind, 1] = fctinter(lag, sz)
        y[ind, 2] = fctinter(splu, sz)
        #        y[ind,3] = fctinter(spfft, sz)

        println("ind=$ind order=$order y[ind,:]=$(y[ind,:])")

        x[ind] = order

        ind += 1
    end

    p = Plots.plot(
        x,
        log10.(y),
        xlabel = "order",
        ylabel = "error",
        legend = :bottomleft,
        label = labels,
        marker = 2,
    )
    prec = precision(BigFloat)

    Plots.savefig(p, "out/resinterp_$(sz)_$(prec)_$(ordmax).pdf")
end

function formule(order, sz, dec)
    return Float64(
        (big(pi) / sz)^(order + 1) * 4 * dec * (1 - dec) / sqrt(big(pi) * (order + 1) / 2),
    )
end

function fctmain_gnuplot(sz, ordmax)
    prec = precision(BigFloat)
    println("# begin gnuplot size=$sz julia precision=$prec")
    println("# order\tLagrange\tformule\t(formule-lag)\tB_SplineLU")
    dec = big"0.351726155665655665187291927162514231451"
    for order = 3:ordmax
        lag = Lagrange(BigFloat, order)
        splu = B_SplineLU(order, sz, big"0.")
        reslag = fctinter(lag, sz, dec)
        ressplu = fctinter(splu, sz, dec)
        resformule = formule(order, sz, dec)
        resdiff = resformule - reslag
        println("$order\t$reslag\t$resformule\t$resdiff\t$ressplu")
    end
end
function fctmainsize_gnuplot(szlist, orderlist)

    prec = precision(BigFloat)
    println("# begin gnuplot order=$orderlist julia precision=$prec")
    print("# size")
    for order in orderlist
        print("\tLagrange($order)\tB_Spline($order)")
    end
    println("")
    dec = big"0.351726155665655665187291927162514231451"
    for sz in szlist
        print("$sz")
        for order in orderlist
            lag = Lagrange(BigFloat, order)
            splu = B_SplineLU(order, sz, big"0.")
            reslag = fctinter(lag, sz, dec)
            ressplu = fctinter(splu, sz, dec)
            print("\t$reslag\t$ressplu")
        end
        println("")
    end
end

setprecision(1024) do
    fctmain_gnuplot(1001, 51)
    #    fctmainsize_gnuplot(101:100:10001,[3,11,31])
end


using SafeTestsets

@safetestset "TensorTrainFactor tests" begin include("TensorTrainFactor.jl") end
@safetestset "TensorTrainFactorization tests" begin include("TensorTrainFactorization.jl") end
@safetestset "FEM tests" begin include("FEM.jl") end


module LinAlgUtils

using SparseArrays
using LinearAlgebra
using StaticArrays
using Rotations

export normalizecolumns
export invertSym3x3
export average_rotmat_l2
export evec_accuracy_report, check_linops, iseigenpairaccurate, eigensystemresiduals
export grassmanian_metric_oper, grassmanian_metric_angle, grassmanian_metric_chordal

# divide all columns of a sparse matrix by their l2 norm.
# Return the new matrix and the vector of column norms
function normalizecolumns(A)
    # get the column sums of A
    S = sqrt.(sum(abs2.(A), dims=1))
    S = reshape(S, (:))

    # get the nonzero entries in A.
    # ei is row index, ej is col index, ev is the value in A
    ei, ej, ev = findnz(A)

    # get the number or rows and columns in A
    m,n = size(A)

    # create a new normalized matrix. For each nonzero index (ei,ej), its new value will be
    # the old value divided by the sum of that column, which can be obtained by S[ej]
    A = sparse(ei,ej,ev./S[ej],m,n)
    return A, S
end

"""
Compute and return the inverse of a Symmetric 3x3 matrix
(this has been carefully profiled / tweaked for performance)
"""
function invertSym3x3(A::Array{Float64,2}) :: SMatrix{3,3, Float64}

    A = SMatrix{3,3}(A) # 6x speedup from this
    detinv = 1 / det(A)

    a = (A[2,2]*A[3,3] - A[2,3]*A[3,2]) * detinv
    b = (A[2,3]*A[3,1] - A[2,1]*A[3,3]) * detinv
    c = (A[2,1]*A[3,2] - A[2,2]*A[3,1]) * detinv

    #d = (A[1,3]*A[3,2] - A[1,2]*A[3,3]) * detinv   # it's symmetric
    e = (A[1,1]*A[3,3] - A[1,3]*A[3,1]) * detinv
    f = (A[1,2]*A[3,1] - A[1,1]*A[3,2]) * detinv

    #g = (A[1,2]*A[2,3] - A[1,3]*A[2,2]) * detinv
    #h = (A[1,3]*A[2,1] - A[1,1]*A[2,3]) * detinv
    i = (A[1,1]*A[2,2] - A[1,2]*A[2,1]) * detinv

    # building array at end: 10% speedup with StaticArrays, slowdown with
    # dynamic memory Arrays.
    return @SMatrix [a b c;
                     b e f;
                     c f i]
end

"""
Compute the average rotation matrix under the l2 geodesic distance. This
algorithm only is guaranteed to work in the rotations lie in a π/2 ball. (This
is generally the case because a π/2 ball covers all but a measure 0 subset of
SO(3).)
"""
function average_rotmat_l2(rots::Vector{RotMatrix{3, Float64}}) :: RotMatrix{3, Float64}
    # Add up all of the rotation matrices using an R^3x3 sum.
    Rsum = SMatrix{3,3}(rots[1])
    for k = 2:length(rots)
        Rsum = Rsum + rots[k]
    end

    # Do an SVD factorization
    fact = svd(Rsum, full=true)

    # form the average matrix
    Ravg = fact.U * fact.Vt

    # correct for cheirality
    if det(Ravg) < 0
        Ravg = fact.U * Diagonal([1.0, 1.0, -1.0]) * fact.Vt
    end

    return Ravg
end

function evec_accuracy_report(operator, evecs, evals; constraints=nothing)

    println("Eigen accuracy report.")
    println("    evecs: ", typeof(evecs), size(evecs))
    N = size(evecs, 2)
    ortho = norm(evecs'*evecs - Matrix{Float64}(I, N, N))
    println("    evecs ortho: |V'V-I|: ", ortho)

    if constraints !== nothing
        Nc = size(constraints, 2)
        println("    constraints: ", typeof(constraints), size(constraints))
        ortho = norm(constraints'*constraints - Matrix{Float64}(I, Nc, Nc))
        println("    constraints ortho: |C'C-I|: ", ortho)
    end

    for (k, (v, λ)) in enumerate(zip(eachcol(evecs), evals))
        v ./= norm(v)
        println("$k. --->")
        # If v is a large eval of covariance
        # it's a small eval of J'J.
        product = operator * v
        λ_actual = norm(product)

        product ./= λ_actual

        println("    stat1: |AX - λX| : ", norm(λ*v - λ_actual * product))
        println("    stat2: |AX| / λ  : ", λ_actual / λ )

        θ = rad2deg(acos(clamp(dot(v, product), -1, 1)))
        println("    stat3: θ deg     : ", θ)
        println("    stat4: dot       : ", dot(v, product))

        println("    supplied λ:    ", λ)
        println("    actual λ=|Av|: ", λ_actual)
        println("    actual λ=v'Av: ", v'*operator*v)

        if constraints !== nothing
            println("    nullspace mass |v'C|: $(norm(v'*constraints))")
        end
    end
end

"""
Perform just a small part of the evec_accuracy_report computation: check the
accuracy of a single eigenvector under the |Av-λv| score. Return True if
the accuracy is satisfactory.
"""
function iseigenpairaccurate(operator, v, λ; ϵ=1e-6) :: Bool
    v ./= norm(v)
    return norm(λ*v - operator*v) < ϵ
end

"""
Return the residuals of each given eigenpair under the |Av-λv| score.
Normalize by a reference scale (Fro norm / num elements) to get a relative measure.
"""
function eigensystemresiduals(operator, evecs, evals) :: Vector{Float64}
    residuals = []
    for (k, (v, λ)) in enumerate(zip(eachcol(evecs), evals))
        v ./= norm(v)
        push!(residuals,  norm(λ*v - operator*v))
    end
    reference = norm(operator, 2) / prod(size(operator))
    return residuals / reference
end

"""
Check if two linear operators are the same.

Allow for them to be in different forms (i.e., dense matrix, SparseMatrixCSC,
    LinearMap, ...). Just assume that they support size() and mult!()

Rather than going element by element (which would take forever, given that we
    don't even have element-wise access) do a randomized search. Generate
    random vectors, apply the operator, and take the norm in the range-space.

Returns the maximum norm-difference across all trials.

Throws if operator dimensions are not identical.
"""
function check_linops(operatorA, operatorB; numtrials=1000)
    # I guess size() doesn't exist for some LinearMaps?
    N = size(operatorA, 1)
    M = size(operatorA, 2)
    # N == size(operatorB, 1) || throw("The operator dimensions are not the same.")
    # M == size(operatorB, 2) || throw("The operator dimensions are not the same.")

    diffs = SharedArray{Float64}(numtrials)
    Threads.@threads for i = 1:numtrials
        x = rand(M)
        diffs[i] = norm(operatorA*x - operatorB*x)
    end
    return maximum(diffs)
end

"""
Measure the distance between two linear subspaces. V and W are both n-by-k
matrices representing elements of the Grassmanian Gr(k, mathbb{R}^n).

Orthogonality is *not* assumed: V represents the subspaces spanned by the
columns of V, and likewise for W.

Return the metric d(V, W) = || P_V - P_W ||
    where P_V is the canonical orthogonal projection onto the column space of V
    where P_W is the canonical orthogonal projection onto the column space of W
    where || ⋅ || is the operator norm
Because the orthogonal projectors are real symmetric, the operator norm is the
largest eigenvalue of the operator.
"""
function grassmanian_metric_oper(V, W)
    size(V) == size(W) || throw("Dimensions don't match!")

    # orthnormalize V and W using a skinny QR
    function orthonormalize(A)
        fact = qr(A)
        return Matrix(fact.Q)
    end
    QV = LinearMap(orthonormalize(V))
    QW = LinearMap(orthonormalize(W))

    # create a matrix-free linear operator for the projection operation
    # PV - PW = (QV*QV') - (QW*QW')
    D = QV * transpose(QV) - QW * transpose(QW)

    # run a matrix-free eigensolver to find the largest eigenvalue
    λ, ϕ, nconv, niter, nmult, resid = eigs(D, nev=1, which=:LM, maxiter=100)
    println("λ: ", λ)
    println("nconv: ", nconv)
    println("niter: ", niter)
    println("nmult: ", nmult)

    return abs(λ[1])
end

"""
Measure the distance between two linear subspaces. V and W are both n-by-k
matrices representing elements of the Grassmanian Gr(k, mathbb{R}^n).

Orthogonality is *not* assumed: V represents the subspace spanned by the
columns of V, and likewise for W.

Return the cosine of the largest principle angle:
        cos θ = max(v in V, w in W, |v|=1, |w|=1)  v'w

Procedure: orthonormalize V and W, then find largest singular value of V'W

Note: this version is a lot faster than the operator norm. I'm seeing about .4ms
    for a 2000x5 problem.
"""
function grassmanian_metric_angle(V, W)
    size(V) == size(W) || throw("Dimensions don't match!")

    # orthnormalize V and W using a skinny QR
    function orthonormalize(A)
        fact = qr(A)
        return Matrix(fact.Q)
    end
    QV = orthonormalize(V)
    QW = orthonormalize(W)

    # assume dims are small: find all singular values with a dense solver.
    # (if #columns is not small, it would be better to use an iterative solver.)
    σ = svdvals(QV' * QW) # returned in descending order
    return σ[1]
end

"""
Measure the distance between two linear subspaces. V and W are both n-by-k
matrices representing elements of the Grassmanian Gr(k, mathbb{R}^n).

Orthogonality is *not* assumed: V represents the subspace spanned by the
columns of V, and likewise for W.

Return the sum of the sins of the principal angles.

Procedure: orthonormalize V and W, then find the singular values of V'W

Citation: Huang, Qui, Calderbank: The Role of Principal Angles in
Subspace Classification
"""
function grassmanian_metric_chordal(V, W)
    size(V) == size(W) || throw("Dimensions don't match!")

    # orthnormalize V and W using a skinny QR
    function orthonormalize(A)
        fact = qr(A)
        return Matrix(fact.Q)
    end
    QV = orthonormalize(V)
    QW = orthonormalize(W)

    # assume dims are small: find all singular values with a dense solver.
    # (if #columns is not small, it would be better to use an iterative solver.)
    σ = svdvals(QV' * QW) # returned in descending order
    # σ's are cos(θ). We want to return ∑sin^2(θ)
    return sum(1 .- σ.^2)
end

end # module


import LinearAlgebra: kron, eigmin

const σx = [0.0 + 0.0im 1; 1 0]
const σy = [0.0 + 0.0im -1.0im; 1.0im 0]
const σz = [1.0 + 0.0im 0; 0 -1]
const σi = [1.0 + 0.0im 0; 0 1]
const σ₋ = [0.0im 0;1 0]
const σ₊ = [0.0im 1;0 0]
const σ = [σx, σy, σz, σi]
xvec = [[1.0 + 0.0im, 1.0] / sqrt(2), [1.0 + 0.0im, -1.0] / sqrt(2)]
yvec = [[1.0, 1.0im] / sqrt(2), [1.0, -1.0im] / sqrt(2)]
zvec = [[1.0 + 0.0im, 0], [0, 1.0 + 0.0im]]

const spσz = sparse(σz)
const spσx = sparse(σx)
const spσy = sparse(σy)
const spσi = sparse(I, 2, 2)

"""
    PauliVec

A constant that holds the eigenvectors of the Pauli matrices. Indices 1, 2 and 3 corresponds to the eigenvectors of ``σ_x``, ``σ_y`` and ``σ_z``.

# Examples
```julia-repl
julia> σx*PauliVec[1][1] == PauliVec[1][1]
true
```
"""
const PauliVec = [xvec, yvec, zvec]

"""
    ⊗(A, B)

Calculate the tensor product of `A` and `B`.

# Examples
```julia-repl
julia> σx⊗σz
4×4 Array{Complex{Float64},2}:
 0.0+0.0im   0.0+0.0im  1.0+0.0im   0.0+0.0im
 0.0+0.0im  -0.0+0.0im  0.0+0.0im  -1.0+0.0im
 1.0+0.0im   0.0+0.0im  0.0+0.0im   0.0+0.0im
 0.0+0.0im  -1.0+0.0im  0.0+0.0im  -0.0+0.0im
```
"""
⊗ = kron

"""
$(SIGNATURES)

Decompse matrix `mat` onto matrix basis `basis`

# Examples
```julia-repl
julia> matrix_decompose(1.0*σx+2.0*σy+3.0*σz, [σx,σy,σz])
3-element Array{Complex{Float64},1}:
 1.0 + 0.0im
 2.0 + 0.0im
 3.0 + 0.0im
```
"""
function matrix_decompose(mat::AbstractMatrix, basis::Vector{<:AbstractMatrix})
    dim = size(basis[1])[1]
    [tr(mat * b) / dim for b in basis]
end

"""
    check_positivity(m)

Check if matrix `m` is positive. Return `true` is the minimum eigenvalue of `m` is greater than or equal to 0.
"""
function check_positivity(m::AbstractMatrix)
    if !ishermitian(m)
        @warn "Input fails the numerical test for Hermitian matrix. Use the upper triangle to construct a new Hermitian matrix."
        d = Hermitian(m)
    else
        d = m
    end
    eigmin(d) >= 0
end

"""
    gibbs_state(h, β)

Calculate the Gibbs state of the matrix `h` at temperature `T` (mK).

# Examples
```julia-repl
julia> gibbs_state(σz, 10)
2×2 Array{Complex{Float64},2}:
 0.178338+0.0im       0.0+0.0im
      0.0+0.0im  0.821662+0.0im
```
"""
function gibbs_state(h, T)
    β = temperature_2_β(T)
    Z = 0.0
    w, v = eigen(Hermitian(h))
    res = zeros(eltype(v), size(h))
    for (i, E) in enumerate(w)
        t = exp(-β * E)
        vi = @view v[:, i]
        res += t * vi * vi'
        Z += t
    end
    res / Z
end

"""
    low_level_matrix(M, lvl)

Calculate the matrix `M` projected to lower energy subspace containing `lvl` energy lvl.

# Examples
```julia-repl
julia> low_level_matrix(σx⊗σx, 2)
4×4 Array{Complex{Float64},2}:
 -0.5+0.0im   0.0+0.0im   0.0+0.0im   0.5+0.0im
  0.0+0.0im  -0.5+0.0im   0.5+0.0im   0.0+0.0im
  0.0+0.0im   0.5+0.0im  -0.5+0.0im   0.0+0.0im
  0.5+0.0im   0.0+0.0im   0.0+0.0im  -0.5+0.0im
```
"""
function low_level_matrix(M, lvl)
    if lvl > size(M, 1)
        @warn "Subspace dimension bigger than total dimension."
        return M
    else
        w, v = eigen(Hermitian(M))
        res = zeros(eltype(v), size(M))
        for i in range(1, stop = lvl)
            vi = @view v[:, i]
            res += w[i] * vi * vi'
        end
        return res
    end
end

"""
$(SIGNATURES)

Test if `𝐔` is a unitary matrix. The function checks how close both `` 𝐔𝐔^† `` and `` 𝐔^†𝐔 `` are to `` I ``, with relative and absolute error given by `rtol`, `atol`.

# Examples
```julia-repl
julia> check_unitary(exp(-1.0im*5*0.5*σx))
true
```
"""
function check_unitary(𝐔::AbstractMatrix; rtol = 1e-6, atol = 1e-8)
    a1 = isapprox(
        𝐔 * 𝐔',
        Matrix{eltype(𝐔)}(I, size(𝐔)),
        rtol = rtol,
        atol = atol,
    )
    a2 = isapprox(
        𝐔' * 𝐔,
        Matrix{eltype(𝐔)}(I, size(𝐔)),
        rtol = rtol,
        atol = atol,
    )
    a1 && a2
end

"""
$(SIGNATURES)

Generate a log-uniformly distributed array with `num` elements between `a` and `b`. The base of log is `base` with default value 10.
"""
function log_uniform(a, b, num; base = 10)
    loga = log(base, a)
    logb = log(base, b)
    10 .^ range(loga, logb, length = num)
end

"""
$(SIGNATURES)

Calculate the partial trace of the density matrix `ρ`, assuming all the subsystems are qubits. `qubit_2_keep` denotes the indices whose corresponding qubits are not traced out.

# Examples
```julia-repl
julia> ρ1 = [0.4 0.2; 0.2 0.6]; ρ2 = [0.5 0; 0 0.5];
julia> partial_trace(ρ1⊗ρ2, [1])
2×2 Array{Float64,2}:
 0.4  0.2
 0.2  0.6
```
"""
function partial_trace(ρ::Matrix, qubit_2_keep::AbstractVector{Int})
    num = Int(log2(size(ρ, 1)))
    partial_trace(ρ, [2 for i in 1:num], qubit_2_keep)
end

"""
$(SIGNATURES)

Calculate the partial trace of the density matrix `ρ`. Assume `ρ` is in the tensor space of ``ℋ₁⊗ℋ₂⊗…``, `sys_dim` is an array of the corresponding sub-system dimensions. `dim_2_keep` is an array of the indices whose corresponding subsystems are not traced out.

# Examples
```julia-repl
julia> ρ1 = [0.4 0.2; 0.2 0.6]; ρ2 = [0.5 0; 0 0.5];
julia> partial_trace(ρ1⊗ρ2, [2, 2], [1])
2×2 Array{Float64,2}:
 0.4  0.2
 0.2  0.6
```
"""
function partial_trace(ρ::Matrix, sys_dim::AbstractVector{Int}, dim_2_keep::AbstractVector{Int})
    size(ρ, 1) != size(ρ, 2) && throw(ArgumentError("ρ is not a square matrix."))
    prod(sys_dim) != size(ρ, 1) && throw(ArgumentError("System dimensions do not multiply to density matrix dimension."))

    N = length(sys_dim)
    sys_dim = reverse(sys_dim)
    ρ = reshape(ρ, sys_dim..., sys_dim...)
    dim_2_keep = N .+ 1 .- dim_2_keep
    dim_2_trace = [i for i in 1:N if !(i in dim_2_keep)]
    traction_idx = collect(1:2*N)
    traction_idx[dim_2_trace .+ N] .= traction_idx[dim_2_trace]
    tensortrace(ρ, traction_idx)
end

"""
$(SIGNATURES)

Calculate the indices of the degenerate values in list `w`. `w` must be sorted in ascending order.

# Examples
```julia-repl
julia> w = [1,1,1,2,3,4,4,5]
julia> find_degenerate(w)
2-element Array{Array{Int64,1},1}:
 [1, 2, 3]
 [6, 7]
```
"""
function find_degenerate(w; atol::Real=1e-6, rtol::Real=0)
    dw = w[2:end] - w[1:end-1]
    inds = findall((x)->isapprox(x, 0.0, atol=atol, rtol=rtol), dw)
    res = Array{Array{Int, 1}, 1}()
    temp = Int[]
    for i in eachindex(inds)
        push!(temp, inds[i])
        if (i==length(inds)) || !(inds[i]+1 == inds[i+1])
            push!(temp, inds[i]+1)
            push!(res, temp)
            temp = Int[]
        end
    end
    res
end

"""
$(SIGNATURES)

Calculate the fidelity between two density matrices `ρ` and `σ`: ``Tr[\\sqrt{\\sqrt{ρ}σ\\sqrt{ρ}}]²``.

# Examples
```julia-repl
julia> ρ = PauliVec[1][1]*PauliVec[1][1]'
julia> σ = PauliVec[3][1]*PauliVec[3][1]'
julia> fidelity(ρ, σ)
0.49999999999999944
```
"""
function fidelity(ρ, σ)
    if !check_density_matrix(ρ) || !check_density_matrix(σ)
        throw(ArgumentError("Inputs are not valid density matrices."))
    end
    temp = sqrt(ρ)
    real(tr(sqrt(temp * σ * temp))^2)
end

"""
$(SIGNATURES)

Check whether matrix `ρ` is a valid density matrix. `atol` and `rtol` is the absolute and relative error tolerance to use when checking `ρ` ≈ 1.

# Examples
```julia-repl
julia> check_density_matrix(σx)
false
julia> check_density_matrix(σi/2)
true
```
"""
function check_density_matrix(ρ; atol::Real=0, rtol::Real=atol>0 ? 0 : √eps())
    ishermitian(ρ) && (eigmin(ρ) >= 0) && isapprox(tr(ρ), 1, atol=atol, rtol=rtol)
end

"""
$(SIGNATURES)

Calculate the purity of density matrix `ρ`: ``tr(ρ²)``. This function does not check whether `ρ` is a valid density matrix.
"""
function purity(ρ)
    tr(ρ*ρ)
end

"""
$(SIGNATURES)

Check whether the input `ρ` is a pure state: ``tr(ρ²)≈1``. This function will first check if `ρ` is a valid density matrix. `atol` and `rtol` are the corresponding absolute and relative error tolerance used for float point number comparison.
"""
function check_pure_state(ρ; atol::Real=0, rtol::Real=atol>0 ? 0 : √eps())
    check_density_matrix(ρ) && isapprox(purity(ρ), 1, atol=atol, rtol=rtol)
end


@add_problem exp function exp_exp_atom(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    y = Variable()
    p = minimize(exp(y), y>=0)
    if test
        @test vexity(p) == ConvexVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 1 atol=atol rtol=rtol
        @test evaluate(exp(y)) ≈ 1 atol=atol rtol=rtol
    end

    y = Variable()
    p = minimize(exp(y), y>=1)
    if test
        @test vexity(p) == ConvexVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ exp(1) atol=atol rtol=rtol
        @test evaluate(exp(y)) ≈ exp(1) atol=atol rtol=rtol
    end

    y = Variable(5)
    p = minimize(sum(exp(y)), y>=0)
    if test
        @test vexity(p) == ConvexVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 5 atol=atol rtol=rtol
        @test evaluate(sum(exp(y))) ≈ 5 atol=atol rtol=rtol
    end

    y = Variable(5)
    p = minimize(sum(exp(y)), y>=0)
    if test
        @test vexity(p) == ConvexVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 5 atol=atol rtol=rtol
    end
end

@add_problem exp function exp_log_atom(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    y = Variable()
    p = maximize(log(y), y<=1)
    if test
        @test vexity(p) == ConvexVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 0 atol=atol rtol=rtol
    end

    y = Variable()
    p = maximize(log(y), y<=2)
    if test
        @test vexity(p) == ConvexVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ log(2) atol=atol rtol=rtol
    end

    y = Variable()
    p = maximize(log(y), [y<=2, exp(y)<=10])
    if test
        @test vexity(p) == ConvexVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ log(2) atol=atol rtol=rtol
    end
end

@add_problem exp function exp_log_sum_exp_atom(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    y = Variable(5)
    p = minimize(logsumexp(y), y>=1)
    if test
        @test vexity(p) == ConvexVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ log(exp(1) * 5) atol=atol rtol=rtol
    end
end

@add_problem exp function exp_logistic_loss_atom(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    y = Variable(5)
    p = minimize(logisticloss(y), y>=1)
    if test
        @test vexity(p) == ConvexVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ log(exp(1) + 1) * 5 atol=atol rtol=rtol
    end
end

@add_problem exp function exp_entropy_atom(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    y = Variable(5, Positive())
    p = maximize(entropy(y), sum(y)<=1)
    if test
        @test vexity(p) == ConvexVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ -(log(1 / 5)) atol=atol rtol=rtol
    end
end

@add_problem exp function exp_relative_entropy_atom(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    x = Variable(1)
    y = Variable(1)
    # x log (x/y)
    p = minimize(relative_entropy(x,y), y==1, x >= 2)
    if test
        @test vexity(p) == ConvexVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 2 * log(2) atol=atol rtol=rtol
    end
end

@add_problem exp function exp_log_perspective_atom(handle_problem!, ::Val{test}, atol, rtol, ::Type{T}) where {T, test}
    x = Variable(1)
    y = Variable(1)
    # y log (x/y)
    p = maximize(log_perspective(x,y), y==5, x <= 10)
    if test
        @test vexity(p) == ConvexVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 5 * log(2) atol=atol rtol=rtol
    end
end


# here we compare various deconvolution options in terms of image quality
using IndexFunArrays, LinearAlgebra, Random, Noise, TestImages, FourierTools
using DeconvOptim, FFTW, Optim, LineSearches
using View5D, Plots

obj = 100f0 .* Float32.(testimage("resolution_test_512"))
# simulate a simple PSF
sz = size(obj);
R_max = sz[1] ./ 12.0;
psf = generate_psf(sz, R_max);

# simulate a perfect image
conv_img = DeconvOptim.conv(obj, psf);

# set a fixed point for the measured data quality
max_photons = 1000

Random.seed!(42)
measured = poisson(conv_img, max_photons);

opt_options = nothing

iterations = 100
function get_data(summary)
    return (summary["best_ncc_img"], summary["best_nvar_img"])
end

function show_ncc!(summary, title="", dpi=300)
        nccs = summary["nccs"]
        plt = plot!(nccs, label=title*" NCC")
        col = plt[1][end].plotattributes[:markercolor]
        vline!([summary["best_ncc_idx"]], line=:dash, color=col, label=title*"_best NCC", dpi=dpi)
        println("best ncc index is $(summary["best_ncc_idx"])")
        println("best ncc is $(summary["best_ncc"])")
        println("best ncc loss is $(summary["losses"][summary["best_ncc_idx"]])")
        println("final ncc is $(summary["nccs"][end])")
        println("final loss is $(summary["losses"][end])")
        xlabel!("iteration")
        ylabel!("normalized cross correlation")
end

function show_nvar!(summary, title="")
    nvars = summary["nvars"]
    nvars_norm = nvars ./ nvars[1]
    plt = plot!(nvars_norm, label=title*" NVAR")
    col = plt[1][end].plotattributes[:markercolor]
    vline!([summary["best_nvar_idx"]], line=:dash, color=col, label=title*"_best NVAR")
    println("best nvar index is $(summary["best_nvar_idx"])")
    println("best nvar is $(summary["best_nvar"])")
    println("best nvar loss is $(summary["losses"][summary["best_nvar_idx"]])")
    println("final nvar is $(summary["nvars"][end])")
    println("final loss is $(summary["losses"][end])")
    xlabel!("iteration")
    ylabel!("normalized variance")
end

function show_loss!(summary, addCurves="", lowest_loss=minimum(summary["losses"]))
        losses = summary["losses"]
        log_losses = log.(losses .- lowest_loss .+ 1)
        rel_losses = log_losses # .- maximum(log_losses)
        plot!(rel_losses[1:end-1], label=addCurves)
        xlabel!("iteration")
        ylabel!("log loss")
end

opt_options, noreg_summary = DeconvOptim.options_trace_deconv(obj, iterations, Non_negative());

res_noreg = deconvolution(measured, psf;
        regularizer=nothing, mapping=Non_negative(),
        opt_options=opt_options, debug_f=nothing)

opt_grad, grad_summary = DeconvOptim.options_trace_deconv(obj, iterations, Non_negative(), );
res_grad = deconvolution(measured, psf;
        regularizer=nothing,  mapping=Non_negative(),
        opt=GradientDescent(), opt_options=opt_grad, debug_f=nothing)

opt_gr, gr_summary = DeconvOptim.options_trace_deconv(obj, iterations, Non_negative());
res_gr = deconvolution(measured, psf;
        regularizer=DeconvOptim.GR(), λ=1e-3,
        mapping=Non_negative(), opt_options=opt_gr, debug_f=nothing)

opt_tv, tv_summary = DeconvOptim.options_trace_deconv(obj, iterations, Non_negative())
res_tv = deconvolution(measured, psf;
        regularizer=DeconvOptim.TV(), λ=1e-3,
        mapping=Non_negative(), opt_options=opt_tv, debug_f=nothing)

plot()
title!("Regularization")
do_display=false
show_ncc!(noreg_summary, "NoReg")
show_nvar!(noreg_summary, "NoReg")
show_ncc!(gr_summary, "GR")
show_nvar!(gr_summary, "GR")
show_ncc!(tv_summary, "TV")
show_nvar!(tv_summary, "TV")

plot()
title!("Optimization")
show_loss!(noreg_summary, "LBFGS", minimum(noreg_summary["losses"]))
show_loss!(grad_summary, "SteepestDecent", minimum(noreg_summary["losses"]))
show_loss!(tv_summary, "TV", minimum(noreg_summary["losses"]))
show_loss!(gr_summary, "GR", minimum(noreg_summary["losses"]))

best_ncc_img, best_nvar_img = get_data(tv_summary)
@vt obj
@vt measured
@vt best_ncc_img
@vt best_nvar_img
@vt res_noreg


## Example:
## DGP - Bivariate VAR(1) Model from Kilian, RESTAT, 1998
# B11 set to 0.01
using VectorAutoregressions, Plots
plotly()

const T,K = 1000,2
const H = 24
const nrep = 1000
const p = 1
const intercept = false
const burnin = 100

srand(1234)
mR  = [0.9 0.01; 0.5 0.5]             #  Coefficients of the 1st lag B11 set to 0.01
mΣ_true = [1 0.3; 0.3 1]              #  vcv of u
mP = cholfact(mΣ_true)[:L]            #  Cholesky PP'= Σ
mI = eye(K)
mIRF_true = At_mul_Bt(mP,mI)[:]       # Initialize structural IRF
# True IRF
for h=1:H
    vIRF = At_mul_Bt(mP,mR^h)[:]
    mIRF_true = hcat(mIRF_true,vIRF)
end

# Generate data from the VAR(1) DGP
y = zeros(T+burnin,K)
gen_var1_data!(y,mR,mP,burnin)

# Fit VAR(1) to data
V = VAR(y,p,intercept)
mIRFb = IRFs_b(V,H,nrep,intercept)
mIRFa = IRFs_a(V,H,intercept)

# Plot true irf vs model irf + asy ci
pIRF_asy = plot(layout = grid(K,K));
[plot!(pIRF_asy, [mIRF_true[i,:] mIRFa.IRF[i,:] mIRFa.CI.CIl[i,:] mIRFa.CI.CIh[i,:]], color = ["blue" "red" "red" "red"],
line = [:solid :solid :dash :dash], legend = false, label = ["true" "model"], subplot = i) for i in 1:K^2]
gui(pIRF_asy)

# Plot true irf vs model irf + bootstraped ci
pIRF_boot = plot(layout = grid(K,K));
[plot!(pIRF_boot, [mIRF_true[i,:] mIRFb.CI.CIl[i,:] mIRFb.IRF[i,:] mIRFb.CI.CIh[i,:]], color = ["blue" "black" "black" "black"],
line = [:solid :solid :dash :dash], legend = false, label = ["true" "model"], subplot = i) for i in 1:K^2]
gui(pIRF_boot)