"""
    countne(x::AbstractUncertainValue, y::AbstractUncertainValue, n::Int)

Count the number of indices at which the elements of two independent length-`n`
draws of `x` and for `y` are not equal.
"""
StatsBase.countne(x::AbstractUncertainValue, y::AbstractUncertainValue, n::Int) =
    resample(StatsBase.countne, x, y, n)

"""
    counteq(x::AbstractUncertainValue, y::AbstractUncertainValue, n::Int)

Count the number of indices at which the elements of two independent length-`n`
draws of `x` and for `y` are equal.
"""
StatsBase.counteq(x::AbstractUncertainValue, y::AbstractUncertainValue, n::Int) =
    resample(StatsBase.counteq, x, y, n)

"""
    corkendall(x::AbstractUncertainValue, y::AbstractUncertainValue, n::Int)

Compute Kendalls's rank correlation coefficient between two uncertain values by
independently drawing `n` samples from `x` and `n` samples from `y`, then computing
Kendalls's rank correlation coefficient between those length-`n` draws.
"""
StatsBase.corkendall(x::AbstractUncertainValue, y::AbstractUncertainValue, n::Int) =
    resample(StatsBase.corspearman, x, y, n)

"""
    corspearman(x::AbstractUncertainValue, y::AbstractUncertainValue, n::Int)

Compute Spearman's rank correlation coefficient between two uncertain values by
independently drawing `n` samples from `x` and `n` samples from `y`, then computing
the Spearman's rank correlation coefficient between those length-`n` draws.
"""
StatsBase.corspearman(x::AbstractUncertainValue, y::AbstractUncertainValue, n::Int) =
    resample(StatsBase.corspearman, x, y, n)

"""
    cor(x::AbstractUncertainValue, y::AbstractUncertainValue, n::Int)

Compute the Pearson correlation between two uncertain values by independently
drawing `n` samples from `x` and `n` samples from `y`, then computing
the Pearson correlation between those length-`n` draws.
"""
StatsBase.cor(x::AbstractUncertainValue, y::AbstractUncertainValue, n::Int) =
    resample(StatsBase.cor, x, y, n)

"""
    cov(x::AbstractUncertainValue, y::AbstractUncertainValue, n::Int;
        corrected::Bool = true)

Compute the covariance between two uncertain values by independently
drawing `n` samples from `x` and `n` samples from `y` , then computing
the covariance between those length-`n` draws.
"""
StatsBase.cov(x::AbstractUncertainValue, y::AbstractUncertainValue, n::Int; corrected::Bool = true) =
    resample(StatsBase.cov, x, y, n; corrected = corrected)

"""
    crosscor(x::AbstractUncertainValue, y::AbstractUncertainValue, [lags],
        n::Int; demean = true)

Compute the cross correlation between two uncertain values by independently
drawing `n` samples from `x` and `n` samples from `y`, `x_draw` and `y_draw`, then computing
the cross correlation between those length-`n` draws. `demean` specifies whether the respective
means of the `x_draw` and `y_draw` should be subtracted from them before computing
their cross correlation.

When left unspecified, the `lags` used are `-min(n-1, 10*log10(n))` to `min(n, 10*log10(n))`.

The output is normalized by `sqrt(var(x_draw)*var(y_draw))`. See `crosscov` for the unnormalized form.
"""
StatsBase.crosscor(x::AbstractUncertainValue, y::AbstractUncertainValue, n::Int; demean = true) =
    resample(StatsBase.crosscor, x, y, n; demean = demean)

StatsBase.crosscor(x::AbstractUncertainValue, y::AbstractUncertainValue, lags, n::Int; demean = true) =
    resample(StatsBase.crosscor, x, y, n, lags; demean = demean)

"""
    crosscov(x::AbstractUncertainValue, y::AbstractUncertainValue, [lags],
        n::Int; demean = true)

Compute the cross covariance function (CCF) between two uncertain values by independently
drawing `n` samples from `x` and `n` samples from `y`, `x_draw` and `y_draw`, then computing
the cross correlation between those length-`n` draws. `demean` specifies whether the respective
means of the `x_draw` and `y_draw` should be subtracted from them before computing
their CCF.

When left unspecified, the `lags` used are `-min(n-1, 10*log10(n))` to `min(n, 10*log10(n))`.

The output is not normalized. See `crosscor` for a function with normalization.
"""
StatsBase.crosscov(x::AbstractUncertainValue, y::AbstractUncertainValue, n::Int; demean = true) =
    resample(StatsBase.crosscov, x, y, n; demean = demean)

StatsBase.crosscov(x::AbstractUncertainValue, y::AbstractUncertainValue, lags, n::Int; demean = true) =
    resample(StatsBase.crosscov, x, y, n, lags; demean = demean)

"""
    gkldiv(x::AbstractUncertainValue, y::AbstractUncertainValue, n::Int)

Compute the generalized Kullback-Leibler divergence between two uncertain
values by independently drawing `n` samples from `x` and `n` samples from `y`,
then computing the generalized Kullback-Leibler divergence between those
length-`n` draws.
"""
StatsBase.gkldiv(x::AbstractUncertainValue, y::AbstractUncertainValue, n::Int) =
    resample(StatsBase.gkldiv, x, y, n)

"""
    kldivergence(x::AbstractUncertainValue, y::AbstractUncertainValue, [b],
        n::Int)

Compute the Kullback-Leibler divergence between two uncertain values by independently
drawing `n` samples from `x` and `n` samples from `y`, then computing the
Kullback-Leibler divergence between those length-`n` draws. Optionally a real number
`b` can be specified such that the divergence is scaled by `1/log(b)`.
"""
StatsBase.kldivergence(x::AbstractUncertainValue, y::AbstractUncertainValue, n::Int) =
    resample(StatsBase.kldivergence, x, y, n)

StatsBase.kldivergence(x::AbstractUncertainValue, y::AbstractUncertainValue, b, n::Int) =
    resample(StatsBase.kldivergence, x, y, n, b)

"""
    maxad(x::AbstractUncertainValue, y::AbstractUncertainValue, n::Int)

Compute the maximum absolute deviation between two uncertain values by independently
drawing `n` samples from `x` and `n` samples from `y`, then computing the
maximum absolute deviation between those length-`n` draws.
"""
StatsBase.maxad(x::AbstractUncertainValue, y::AbstractUncertainValue, n::Int) =
    resample(StatsBase.maxad, x, y, n)

"""
    meanad(x::AbstractUncertainValue, y::AbstractUncertainValue, n::Int)

Compute the mean absolute deviation between two uncertain values by independently
drawing `n` samples from `x` and `n` samples from `y`, then computing the
mean absolute deviation between those length-`n` draws.
"""
StatsBase.meanad(x::AbstractUncertainValue, y::AbstractUncertainValue, n::Int) =
    resample(StatsBase.meanad, x, y, n)

"""
    msd(x::AbstractUncertainValue, y::AbstractUncertainValue, n::Int)

Compute the mean squared deviation between two uncertain values by independently
drawing `n` samples from `x` and `n` samples from `y`, then computing the
mean squared deviation between those length-`n` draws.
"""
StatsBase.msd(x::AbstractUncertainValue, y::AbstractUncertainValue, n::Int) =
    resample(StatsBase.msd, x, y, n)

"""
    psnr(x::AbstractUncertainValue, y::AbstractUncertainValue, maxv, n::Int)

Compute the peak signal-to-noise ratio between two uncertain values by independently
drawing `n` samples from `x` and from `y`, yielding `x_draw` and `y_draw`, then
computing the peak signal-to-noise ratio between those length-`n` draws.

The PSNR is computed as `10 * log10(maxv^2 / msd(x_draw, y_draw))`, where `maxv` is
the maximum possible value `x` or `y` can take
"""
StatsBase.psnr(x::AbstractUncertainValue, y::AbstractUncertainValue, maxv, n::Int) =
    resample(StatsBase.psnr, x, y, n, maxv)

"""
    rmsd(x::AbstractUncertainValue, y::AbstractUncertainValue, n::Int,
        normalize = false)

Compute the root mean squared deviation between two uncertain values by independently
drawing `n` samples from `x` and from `y`, yielding `x_draw` and `y_draw`, then
computing the the root mean squared deviation between those length-`n` draws.
The root mean squared deviation is computed as `sqrt(msd(x_draw, y_draw))`.
Optionally, `x_draw` and `y_draw` may be normalised.
"""
StatsBase.rmsd(x::AbstractUncertainValue, y::AbstractUncertainValue, n::Int; normalize = false) =
    resample(StatsBase.rmsd, x, y, n; normalize = normalize)

"""
    sqL2dist(x::AbstractUncertainValue, y::AbstractUncertainValue, n::Int)

Compute the squared L2 distance between two uncertain values by independently
drawing `n` samples from `x` and from `y`, then computing the
squared L2 distance between those length-`n` draws: ``\\sum_{i=1}^n |x_i - y_i|^2``.
"""
StatsBase.sqL2dist(x::AbstractUncertainValue, y::AbstractUncertainValue, n::Int) =
    resample(StatsBase.sqL2dist, x, y, n)


@testset "Random spatial vector" begin
    for i = 1:10
        @test norm(rand(SpatialVector{Float64})) ≈ 1.
    end;
end;


using Test, PolyChaos
import LinearAlgebra: norm, dot

##########################################################
# MEASURE
##########################################################
@testset "Error handling for Measure" begin
    @test_throws AssertionError Measure("jacobi",Dict(:shape_a=>-1.,:shape_b=>2.))
    @test_throws AssertionError Measure("jacobi",Dict(:shape_a=>-1.,:shape_b=>-2.))
    @test_throws AssertionError Measure("jacobi",Dict(:shape_a=>1,:shape_b=>-2))

    @test_throws AssertionError Measure("beta01",Dict(:shape_a=>0.,:shape_b=>2.))
    @test_throws AssertionError Measure("beta01",Dict(:shape_a=>0.,:shape_b=>-2.))
    @test_throws AssertionError Measure("beta01",Dict(:shape_a=>1,:shape_b=>-2))

    @test_throws ErrorException Measure("something")

    @test_throws AssertionError Measure("test",t->t,(3,-1),true)
end

##########################################################
# ORTHOPOLY
##########################################################
@test_throws ErrorException OrthoPoly("idonotexist",3)

quads = [ fejer, fejer2, clenshaw_curtis ]
discs = [ lanczos, stieltjes ]
degs  = 1:5:20
m = Measure("uniform01")
tol = 1e-6
@testset "General constructor for Gaussian OrthoPoly" begin
    for quad in quads, disc in discs, deg in degs
        op = OrthoPoly("uniform01",deg)
        op1 = OrthoPoly("myuniform",deg,m; quadrature=quad, discretization = disc, Nquad = 100*deg)
        @test isapprox(norm(coeffs(op)-coeffs(op1),Inf),0; atol = tol)
        op2 = OrthoPoly("myuniform",deg,t->1,(0,1),true; quadrature=quad, discretization = disc, Nquad = 100*deg)
        @test isapprox(norm(coeffs(op)-coeffs(op2),Inf),0; atol = tol)
    end
end

##########################################################
# QUAD
##########################################################
Ns = 10:5:30
m = Measure("uniform01")
myf(t) = 6t^5
result = 1.
tol = 1e-8
@testset "Last-resort-constructor for Quad using quadgp" begin
    for quad in quads, N in Ns
        op = OrthoPoly("myOP",N,m)
        NW = nw(Quad(N,m;quadrature=quad))
        @test isapprox(dot(myf.(NW[:,1]),NW[:,2]), result; atol=tol)
        NW = nw(Quad(N,t->1.,(0,1),true;quadrature=quad))
        @test isapprox(dot(myf.(NW[:,1]),NW[:,2]), result; atol=tol)

        NW = nw(Quad(N,op))
        @test isapprox(dot(myf.(NW[:,1]),NW[:,2]), result; atol=tol)
        NW = nw(Quad(N,m.w,op.α,op.β,m.dom,m.symmetric,m.pars))
        @test isapprox(dot(myf.(NW[:,1]),NW[:,2]), result; atol=tol)
    end
end

d = 5
ops = [ OrthoPoly("hermite",d),
        OrthoPoly("uniform01",d),
        OrthoPoly("beta01",d,Dict(:shape_a=>3.2,:shape_b=>5.34)) ]
mop = MultiOrthoPoly(ops,d)

is = 0:d
n = length(ops)

x = [1., 2., 3.]

@testset "Evaluation of multivariate basis" begin
    for ind in Iterators.product([collect(0:d) for i = 1:n]...)
        @test isapprox(prod( map(i->evaluate(ind[i],x[i],ops[i]),1:n) ) -  evaluate([ind...],x,mop)[1],0;atol=tol)
    end
end

##########################################################
# TENSOR
##########################################################
deg, M = 5, 1:4
opq = OrthoPolyQ("legendre",deg;Nrec=3*deg)
mop = MultiOrthoPoly([OrthoPolyQ("gaussian",deg;Nrec=3*deg), OrthoPolyQ("logistic",deg;Nrec=3*deg), OrthoPolyQ("uniform01",deg;Nrec=3*deg)], deg)

@test_throws AssertionError Tensor(0,mop)
@testset "Check tensor" begin
    for m in M, op in [opq, mop]
        tensor = Tensor(m,op)
        for ind in Iterators.product([collect(0:deg) for i=1:m]...)
            @test isapprox(tensor.get(collect(ind)), computeSP(collect(ind),op); atol = tol )
        end
    end
end


export moment, moment_value, monomial, base_function

struct Moment{T, PT <: AbstractPolynomialLike} <: AbstractMoment{T}
    α::T
    x::PT
end

"""
    moment(α, m::AbstractPolynomialLike)

Creates the moment of the base function `m` of value `α`.
"""
moment(α, m::AbstractPolynomialLike) = Moment(α, m)

"""
    moment_value(m::AbstractMomentLike)

Returns the value of the moment `m`.

## Examples

Calling `moment_value(moment(3.1, x*y^2))` should return `3.1`.
"""
moment_value(m::Moment) = m.α

"""
    monomial(m::AbstractMomentLike)

Returns the monomial of the moment `m`.

## Examples

Calling `monomial(moment(3.1, x*y^2))` should return `x*y^2`.
"""
MP.monomial(m::Moment) = m.x

"""
    base_function(m::AbstractMomentLike)

Returns the base_function of the moment `m`.

## Examples

Calling `base_function(moment(3.1, x*y^2))` should return `x*y^2`.
"""
base_function(m::Moment) = m.x

for f in [:variables, :nvariables, :exponents, :powers]
    @eval begin
        MP.$f(m::AbstractMomentLike) = $f(base_function(m))
    end
end

MP.degree(m::AbstractMomentLike) = MP.maxdegree(base_function(m))


using FlightSims
import FlightSims: State, Dynamics!
using DataFrames
using ComponentArrays
using UnPack
using Transducers
using Plots
using SciMLBase
using Test

struct LinearSystem_ZOH_Input <: AbstractEnv
    linear_env::LinearSystem
end

function State(env::LinearSystem_ZOH_Input)
    State(env.linear_env)
end

function Dynamics!(env::LinearSystem_ZOH_Input)
    @Loggable function dynamics!(dx, x, input, t)
        @nested_log Dynamics!(env.linear_env)(dx, x, nothing, t; u=input)
    end
end

function main()
    A = [0 1;
         0 0]
    B = [0 1]'
    linear_env = LinearSystem(A, B)
    env = LinearSystem_ZOH_Input(linear_env)
    x0_state = [1, 2]
    input = zeros(1)
    x0 = State(env)(x0_state)
    t0 = 0.0
    tf = 20.0
    simulator = Simulator(x0, Dynamics!(env), input; tf=tf)
    # interactive sim
    Δt = 0.1  # save period
    df = DataFrame()
    @time for (i, t) in enumerate(t0:Δt:tf)
        # To perform interactive simulation,
        # you should be aware of the integrator interface
        # provided by DifferentialEquations.jl;
        # see https://diffeq.sciml.ai/stable/basics/integrator/#integrator
        # e.g., DO NOT directly change the integrator state
        state = simulator.integrator.u
        input = simulator.integrator.p
        if (i-1) % 10 == 0  # update input period
            input .= -sum(state)
        end
        step_until!(simulator, t, df)
    end
    # plot
    ts = df.time
    states = df.sol |> Map(datum -> datum.state) |> collect
    inputs = df.sol |> Map(datum -> datum.input) |> collect
    p_x = plot(ts, hcat(states...)';
               title="state variable", label=["x1" "x2"],
              )
    p_u = plot(ts, hcat(inputs...)';
               title="input variable", label="u",
               linetype=:steppost,  # to plot zero-order-hold input appropriately
              )
    fig = plot(p_x, p_u; layout=(2, 1))
    savefig("figures/interactive_sim.png")
    display(fig)
end

@testset "linear_system_zoh_input" begin
    main()
end


# See p. 280 of Murphy's Machine Learning
# x_k1 = x_k - alpha * gr + mu * (x - x_previous)

struct MomentumGradientDescent{Tf, IL,L} <: FirstOrderOptimizer
    mu::Tf
    alphaguess!::IL
    linesearch!::L
    manifold::Manifold
end

Base.summary(::MomentumGradientDescent) = "Momentum Gradient Descent"

function MomentumGradientDescent(; mu::Real = 0.01,
                                 alphaguess = LineSearches.InitialPrevious(), # TODO: investigate good defaults
                                 linesearch = LineSearches.HagerZhang(),        # TODO: investigate good defaults
                                 manifold::Manifold=Flat())
    MomentumGradientDescent(mu, _alphaguess(alphaguess), linesearch, manifold)
end

mutable struct MomentumGradientDescentState{Tx, T} <: AbstractOptimizerState
    x::Tx
    x_previous::Tx
    x_momentum::Tx
    f_x_previous::T
    s::Tx
    @add_linesearch_fields()
end

function initial_state(method::MomentumGradientDescent, options, d, initial_x)
    T = eltype(initial_x)
    initial_x = copy(initial_x)
    retract!(method.manifold, initial_x)

    value_gradient!!(d, initial_x)

    project_tangent!(method.manifold, gradient(d), initial_x)

    MomentumGradientDescentState(initial_x, # Maintain current state in state.x
                                 copy(initial_x), # Maintain previous state in state.x_previous
                                 similar(initial_x), # Record momentum correction direction in state.x_momentum
                                 real(T)(NaN), # Store previous f in state.f_x_previous
                                 similar(initial_x), # Maintain current search direction in state.s
                                 @initial_linesearch()...)
end

function update_state!(d, state::MomentumGradientDescentState, method::MomentumGradientDescent)
    project_tangent!(method.manifold, gradient(d), state.x)
    # Search direction is always the negative gradient
    state.s .= .-gradient(d)

    # Update position, and backup current one
    state.x_momentum .= state.x .- state.x_previous

    # Determine the distance of movement along the search line
    lssuccess = perform_linesearch!(state, method, ManifoldObjective(method.manifold, d))

    state.x .+= state.alpha.*state.s .+ method.mu.*state.x_momentum
    retract!(method.manifold, state.x)
    lssuccess == false # break on linesearch error
end

function trace!(tr, d, state, iteration, method::MomentumGradientDescent, options, curr_time=time())
  common_trace!(tr, d, state, iteration, method, options, curr_time)
end

function default_options(method::MomentumGradientDescent)
    Dict(:allow_f_increases => true)
end


mutable struct QpJuMP{T,Tv<:AbstractArray{T},Tm<:AbstractMatrix{T}} <: AbstractSubOptimizer
    model::JuMP.Model
    data::QpData{T,Tv,Tm}
    x::Vector{JuMP.VariableRef}
    constr::Vector{JuMP.ConstraintRef}
    rngbdcons::Vector{Int}
    rngcons::Vector{Int}
    slack_vars::Dict{Int,Vector{JuMP.VariableRef}}

    function QpJuMP(model::JuMP.AbstractModel, data::QpData{T,Tv,Tm}) where {T,Tv,Tm}
        qp = new{T,Tv,Tm}()
        qp.model = model
        qp.data = data
        qp.x = []
        qp.constr = []
        qp.rngbdcons = []
        qp.rngcons = []
        qp.slack_vars = Dict()
        return qp
    end
end

SubOptimizer(model::JuMP.AbstractModel, data::QpData{T,Tv,Tm}) where {T,Tv,Tm} =
    QpJuMP(model, data)

"""
    create_model!

Initialize QP subproblem in JuMP.Model. The model assumes that the first `qp.data.num_linear_constraints` constraints are linear.
The slack variables are not introduced for the linear constraints.

# Arguments
- `qp`
- `Δ`: trust-region size
"""
function create_model!(
    qp::QpJuMP{T,Tv,Tm},
    Δ::T,
) where {T,Tv,Tm}

    qp.constr = []
    qp.rngbdcons = []
    qp.rngcons = []
    empty!(qp.slack_vars)

    n = length(qp.data.c)
    m = length(qp.data.c_lb)

    # create nominal variables
    qp.x = @variable(
        qp.model,
        [i = 1:n],
        base_name = "x",
    )

    set_trust_region!(qp, Δ)

    # add slack variables only for nonlinear constraints
    for i = (qp.data.num_linear_constraints+1):m
        qp.slack_vars[i] = []
        push!(qp.slack_vars[i], @variable(qp.model, base_name = "u[$i]", lower_bound = 0.0))
        if qp.data.c_lb[i] > -Inf && qp.data.c_ub[i] < Inf
            push!(qp.slack_vars[i], @variable(qp.model, base_name = "v[$i]", lower_bound = 0.0))
        end
    end

    # dummy objective function
    @objective(qp.model, Min, 0.0)

    # create affine constraints
    for i = 1:m
        c_ub = qp.data.c_ub[i]
        c_lb = qp.data.c_lb[i]
        if abs(qp.data.b[i]) < Inf
            c_ub -= qp.data.b[i]
            c_lb -= qp.data.b[i]
        end

        if qp.data.c_lb[i] == qp.data.c_ub[i] #This means the constraint is equality
            if i <= qp.data.num_linear_constraints
                Arow = qp.data.A[i,:]
                push!(
                    qp.constr,
                    @constraint(qp.model, sum(v * qp.x[Arow.nzind[j]] for (j,v) in enumerate(Arow.nzval)) == c_lb)
                )
            else
                push!(qp.constr, @constraint(qp.model, qp.slack_vars[i][1] - qp.slack_vars[i][2] == c_lb))
            end
        elseif qp.data.c_lb[i] > -Inf && qp.data.c_ub[i] < Inf
            if i <= qp.data.num_linear_constraints
                Arow = qp.data.A[i,:]
                # push!(qp.constr, @constraint(qp.model, c_lb <= sum(A[i,j] * qp.x[j] for j in A[i,:].nzind) <= c_ub))
                push!(qp.constr, @constraint(qp.model, sum(v * qp.x[Arow.nzind[j]] for (j,v) in enumerate(Arow.nzval)) >= c_lb))
            else
                push!(qp.constr, @constraint(qp.model, qp.slack_vars[i][1] >= c_lb))
            end
            push!(qp.rngcons, i)
        elseif qp.data.c_lb[i] > -Inf
            if i <= qp.data.num_linear_constraints
                Arow = qp.data.A[i,:]
                push!(qp.constr, @constraint(qp.model, sum(v * qp.x[Arow.nzind[j]] for (j,v) in enumerate(Arow.nzval)) >= c_lb))
            else
                push!(qp.constr, @constraint(qp.model, qp.slack_vars[i][1] >= c_lb))
            end
        elseif qp.data.c_ub[i] < Inf
            if i <= qp.data.num_linear_constraints
                Arow = qp.data.A[i,:]
                push!(qp.constr, @constraint(qp.model, sum(v * qp.x[Arow.nzind[j]] for (j,v) in enumerate(Arow.nzval)) <= c_ub))
            else
                push!(qp.constr, @constraint(qp.model, -qp.slack_vars[i][1] <= c_ub))
            end
        end
    end

    # create ranged affine constraints
    for i in qp.rngcons
        c_ub = qp.data.c_ub[i] - qp.data.b[i]
        if i <= qp.data.num_linear_constraints
            Arow = qp.data.A[i,:]
            push!(qp.constr, @constraint(qp.model, sum(v * qp.x[Arow.nzind[j]] for (j,v) in enumerate(Arow.nzval)) <= c_ub))
        else
            push!(qp.constr, @constraint(qp.model, -qp.slack_vars[i][2] <= c_ub))
        end
    end
end

function sub_optimize!(
    qp::QpJuMP{T,Tv,Tm},
    x_k::Tv,
    Δ::T,
) where {T,Tv,Tm}

    # dimension of LP
    m, n = size(qp.data.A)

    # modify objective function
    if isnothing(qp.data.Q)
        @objective(qp.model, qp.data.sense,
            sum(qp.data.c[i] * qp.x[i] for i = 1:n)
        )
    else
        # obj = QuadExpr(
        #     sum(qp.data.c[i] * qp.x[i] for i = 1:n)
        # )
        # for j = 1:qp.data.Q.n, i in nzrange(qp.data.Q, j)
        #     add_to_expression!(
        #         obj,
        #         0.5*qp.data.Q.nzval[i],
        #         qp.x[qp.data.Q.rowval[i]],
        #         qp.x[j],
        #     )
        # end
        # @objective(qp.model, qp.data.sense, obj)
        @objective(
            qp.model,
            qp.data.sense,
            sum(qp.data.c[i] * qp.x[i] for i = 1:n)
            + 0.5 * sum(
                qp.data.Q.nzval[i] * qp.x[qp.data.Q.rowval[i]] * qp.x[j]
                for j = 1:qp.data.Q.n for i in nzrange(qp.data.Q, j)
            )
        )
    end

    # fix slack variables to zeros
    for (_, slacks) in qp.slack_vars, s in slacks
        if JuMP.has_lower_bound(s)
            JuMP.delete_lower_bound(s)
        end
        JuMP.fix(s, 0.0)
    end

    set_trust_region!(qp, x_k, Δ)
    modify_constraints!(qp)

    # @show x_k
    # JuMP.print(qp.model)
    JuMP.optimize!(qp.model)
    status = termination_status(qp.model)
    Xsol, lambda, mult_x_U, mult_x_L, p_slack = collect_solution!(qp, status)

    return Xsol, lambda, mult_x_U, mult_x_L, p_slack, status
end

function sub_optimize_lp(
    optimizer,
    A::Tm, cl::Tv, cu::Tv,
    xl::Tv, xu::Tv, x_k::Tv,
    m::Int, num_constraints::Int
) where {T, Tv<:AbstractArray{T}, Tm<:AbstractMatrix{T}}
    n = length(x_k)

    model = JuMP.Model(optimizer)
    @variable(model, xl[i] <= x[i=1:n] <= xu[i])
    @objective(model, Min, sum((x[i] - x_k[i])^2 for i=1:n))
    constr = Vector{JuMP.ConstraintRef}(undef, m)
    for i = 1:m
        arow = A[i,:]
        if cl[i] == cu[i]
            constr[i] = @constraint(model, sum(a * x[arow.nzind[j]] for (j, a) in enumerate(arow.nzval)) == cl[i])
        elseif cl[i] > -Inf && cu[i] < Inf
            constr[i] = @constraint(model, cl[i] <= sum(a * x[arow.nzind[j]] for (j, a) in enumerate(arow.nzval)) <= cu[i])
        elseif cl[i] > -Inf
            constr[i] = @constraint(model, sum(a * x[arow.nzind[j]] for (j, a) in enumerate(arow.nzval)) >= cl[i])
        elseif cu[i] < Inf
            constr[i] = @constraint(model, sum(a * x[arow.nzind[j]] for (j, a) in enumerate(arow.nzval)) <= cu[i])
        end
    end
    JuMP.optimize!(model)
    status = termination_status(model)

    Xsol = Tv(undef, n)
    lambda = zeros(T, num_constraints)
    mult_x_U = zeros(T, n)
    mult_x_L = zeros(T, n)

    if status ∈ [MOI.OPTIMAL, MOI.ALMOST_OPTIMAL, MOI.ALMOST_LOCALLY_SOLVED, MOI.LOCALLY_SOLVED]
        Xsol .= JuMP.value.(x)

        # extract the multipliers to constraints
        for i = 1:m
            lambda[i] = JuMP.dual(constr[i])
        end

        # extract the multipliers to column bounds
        for i = 1:n
            redcost = JuMP.reduced_cost(x[i])
            if redcost > 0
                mult_x_L[i] = redcost
            elseif redcost < 0
                mult_x_U[i] = redcost
            end
        end
    elseif status ∈ [MOI.LOCALLY_INFEASIBLE, MOI.INFEASIBLE, MOI.DUAL_INFEASIBLE, MOI.NORM_LIMIT, MOI.OBJECTIVE_LIMIT]
        fill!(Xsol, 0.0)
        fill!(lambda, 0.0)
        fill!(mult_x_U, 0.0)
        fill!(mult_x_L, 0.0)
    else
        @error "Unexpected status: $(status)"
    end

    return Xsol, lambda, mult_x_U, mult_x_L, status
end

function sub_optimize_lp(
    qp::QpJuMP{T,Tv,Tm},
    x_k::Tv,
) where {T,Tv,Tm}

    # problem dimension
    m, n = size(qp.data.A)

    # modify objective function
    @objective(qp.model, Min, sum((qp.x[i] - x_k[i])^2 for i=1:n))
    # @objective(qp.model, Min, sum((qp.x[i])^2 for i=1:n))

    # modify slack variable bounds
    for (i, slacks) in qp.slack_vars, s in slacks
        if JuMP.has_lower_bound(s)
            JuMP.delete_lower_bound(s)
        end
        if i <= qp.data.num_linear_constraints
            JuMP.fix(s, 0.0)
        end
    end

    # set initial variable values
    for i = 1:n
        JuMP.set_start_value(qp.x[i], x_k[i])
    end

    set_trust_region!(qp, Inf)
    modify_constraints!(qp)

    JuMP.optimize!(qp.model)
    status = termination_status(qp.model)
    Xsol, lambda, mult_x_U, mult_x_L, p_slack = collect_solution!(qp, status)

    return Xsol, lambda, mult_x_U, mult_x_L, status
end

function sub_optimize_L1QP!(
    qp::QpJuMP{T,Tv,Tm},
    x_k::Tv,
    Δ::T,
    μ::T,
) where {T,Tv,Tm}

    # problem dimension
    m, n = size(qp.data.A)

    # modify objective function
    obj_direction = ifelse(qp.data.sense == MOI.MIN_SENSE, 1.0, -1.0)
    if isnothing(qp.data.Q)
        @objective(qp.model, qp.data.sense,
            sum(qp.data.c[i] * qp.x[i] for i = 1:n)
            + obj_direction * μ * sum(s for (_, slacks) in qp.slack_vars, s in slacks)
        )
    else
        # μ_inv = 1.0 / μ
        # @objective(
        #     qp.model,
        #     qp.data.sense,
        #     μ_inv * sum(qp.data.c[i] * qp.x[i] for i = 1:n)
        #     + μ_inv * 0.5 * sum(
        #         qp.data.Q.nzval[i] * qp.x[qp.data.Q.rowval[i]] * qp.x[j]
        #         for j = 1:qp.data.Q.n for i in nzrange(qp.data.Q, j)
        #     )
        #     + obj_direction * sum(s for (_, slacks) in qp.slack_vars for s in slacks)
        # )
        @objective(
            qp.model,
            qp.data.sense,
            sum(qp.data.c[i] * qp.x[i] for i = 1:n)
            + 0.5 * sum(
                qp.data.Q.nzval[i] * qp.x[qp.data.Q.rowval[i]] * qp.x[j]
                for j = 1:qp.data.Q.n for i in nzrange(qp.data.Q, j)
            )
            + obj_direction * μ * sum(s for (_, slacks) in qp.slack_vars for s in slacks)
        )
    end

    # modify slack variable bounds
    for (_, slacks) in qp.slack_vars, s in slacks
        if JuMP.is_fixed(s)
            JuMP.unfix(s)
        end
        set_lower_bound(s, 0.0)
    end

    set_trust_region!(qp, x_k, Δ)
    modify_constraints!(qp)

    # JuMP.print(qp.model)
    JuMP.optimize!(qp.model)
    status = termination_status(qp.model)
    Xsol, lambda, mult_x_U, mult_x_L, p_slack = collect_solution!(qp, status)

    # for i in 1:n
    #     v_lb = qp.data.v_lb[i] - x_k[i]
    #     v_ub = qp.data.v_ub[i] - x_k[i]
    #     @show i, Xsol[i], v_lb, v_ub, mult_x_L[i], mult_x_U[i]
    # end

    return Xsol, lambda, mult_x_U, mult_x_L, p_slack, status
end

"""
Solve QP subproblem for feasibility restoration
"""
function sub_optimize_FR!(
    qp::QpJuMP{T,Tv,Tm},
    x_k::Tv,
    Δ::T,
) where {T,Tv,Tm}

    # dimension of LP
    m, n = size(qp.data.A)

    # modify objective function
    @objective(qp.model, Min, sum(s for (_, slacks) in qp.slack_vars, s in slacks))

    # modify slack variable bounds
    for (i, slacks) in qp.slack_vars
        if qp.data.b[i] >= qp.data.c_lb[i] && qp.data.b[i] <= qp.data.c_ub[i]
            for s in slacks
                if JuMP.is_fixed(s) == false
                    JuMP.fix(s, 0.0, force = true)
                end
            end
        else
            for s in slacks
                if JuMP.is_fixed(s)
                    JuMP.unfix(s)
                end
                set_lower_bound(s, 0.0)
            end
        end
    end

    set_trust_region!(qp, x_k, Δ)
    modify_constraints!(qp)

    # JuMP.write_to_file(qp.model, "debug_jump.lp", format = MOI.FileFormats.FORMAT_LP)
    # @show x_k
    # JuMP.print(qp.model)
    JuMP.optimize!(qp.model)
    status = termination_status(qp.model)
    Xsol, lambda, mult_x_U, mult_x_L, p_slack = collect_solution!(qp, status)

    return Xsol, lambda, mult_x_U, mult_x_L, p_slack, status
end

"""
Compute the infeasibility of the linearized model
"""
function sub_optimize_infeas(
    qp::QpJuMP{T,Tv,Tm},
    x_k::Tv,
    Δ::T,
) where {T,Tv,Tm}

    # modify objective function
    @objective(qp.model, Min, sum(s for (_, slacks) in qp.slack_vars, s in slacks))

    # modify slack variable bounds
    for (_, slacks) in qp.slack_vars, s in slacks
        if JuMP.is_fixed(s)
            JuMP.unfix(s)
        end
        set_lower_bound(s, 0.0)
    end

    set_trust_region!(qp, x_k, Δ)
    modify_constraints!(qp)

    JuMP.optimize!(qp.model)
    status = termination_status(qp.model)

    Xsol = Tv(undef, length(qp.x))
    infeasibility = Inf
    if status ∈ [MOI.OPTIMAL, MOI.ALMOST_LOCALLY_SOLVED, MOI.LOCALLY_SOLVED]
        Xsol .= JuMP.value.(qp.x)
        infeasibility = JuMP.objective_value(qp.model)
    end

    return Xsol, infeasibility
end

function set_trust_region!(
    x::Vector{JuMP.VariableRef},
    v_lb::Tv,
    v_ub::Tv,
    Δ::T
) where {T,Tv}
    for i in eachindex(x)
        set_lower_bound(x[i], max(-Δ, v_lb[i]))
        set_upper_bound(x[i], min(+Δ, v_ub[i]))
    end
end

function set_trust_region!(
    qp::QpJuMP{T,Tv,Tm},
    x_k::Tv,
    Δ::T
) where {T,Tv,Tm}
    return set_trust_region!(qp.x, qp.data.v_lb - x_k, qp.data.v_ub - x_k, Δ)
end

function set_trust_region!(
    qp::QpJuMP{T,Tv,Tm},
    Δ::T
) where {T,Tv,Tm}
    return set_trust_region!(qp.x, qp.data.v_lb, qp.data.v_ub, Δ)
end

function modify_constraints!(qp::QpJuMP{T,Tv,Tm}) where {T,Tv,Tm}

    # problem dimension
    m, n = size(qp.data.A)

    # modify the nonlinear constraint coefficients
    for j = 1:qp.data.A.n, i in nzrange(qp.data.A, j)
        if qp.data.A.rowval[i] > qp.data.num_linear_constraints
            set_normalized_coefficient(
                qp.constr[qp.data.A.rowval[i]],
                qp.x[j],
                qp.data.A.nzval[i],
            )
        end
    end

    # modify the coefficients for the other part of ranged constraints
    for (ind, val) in enumerate(qp.rngcons)
        if val > qp.data.num_linear_constraints
            row_of_A = qp.data.A[val, :]
            for (i,j) = enumerate(row_of_A.nzind)
                set_normalized_coefficient(qp.constr[m+ind], qp.x[j], row_of_A.nzval[i])
            end
        end
    end

    # modify RHS
    for i in 1:m
        c_ub = qp.data.c_ub[i] - qp.data.b[i]
        c_lb = qp.data.c_lb[i] - qp.data.b[i]

        if qp.data.c_lb[i] == qp.data.c_ub[i]
            set_normalized_rhs(qp.constr[i], c_lb)
        elseif qp.data.c_lb[i] > -Inf && qp.data.c_ub[i] < Inf
            set_normalized_rhs(qp.constr[i], c_lb)
        elseif qp.data.c_lb[i] > -Inf
            set_normalized_rhs(qp.constr[i], c_lb)
        elseif qp.data.c_ub[i] < Inf
            set_normalized_rhs(qp.constr[i], c_ub)
        end
    end

    # modify the RHS for the other part of ranged constraints
    for (i, val) in enumerate(qp.rngcons)
        c_ub = qp.data.c_ub[val] - qp.data.b[val]
        set_normalized_rhs(qp.constr[i+m], c_ub)
    end
end

function collect_solution!(qp::QpJuMP{T,Tv,Tm}, status) where {T,Tv,Tm}

    # problem dimension
    m, n = size(qp.data.A)

    Xsol = Tv(undef, n)
    p_slack = Dict{Int,Vector{Float64}}()
    lambda = Tv(undef, m)
    mult_x_U = zeros(T, n)
    mult_x_L = zeros(T, n)

    if status ∈ [MOI.OPTIMAL, MOI.ALMOST_OPTIMAL, MOI.ALMOST_LOCALLY_SOLVED, MOI.LOCALLY_SOLVED]
        Xsol .= JuMP.value.(qp.x)
        for (i, slacks) in qp.slack_vars
            p_slack[i] = JuMP.value.(slacks)
        end
        # @show JuMP.objective_value(qp.model), Xsol
        # @show p_slack

        # extract the multipliers to constraints
        for i = 1:m
            lambda[i] = JuMP.dual(qp.constr[i])
        end
        for (i, val) in enumerate(qp.rngcons)
            lambda[val] += JuMP.dual(qp.constr[i+m])
        end
        # @show MOI.get(qp.model, MOI.ConstraintDual(1), qp.constr)

        # extract the multipliers to column bounds
        for i = 1:n
            redcost = JuMP.reduced_cost(qp.x[i])
            if redcost > 0
                mult_x_L[i] = redcost
            elseif redcost < 0
                mult_x_U[i] = redcost
            end
        end
    elseif status ∈ [MOI.LOCALLY_INFEASIBLE, MOI.INFEASIBLE, MOI.DUAL_INFEASIBLE, MOI.NORM_LIMIT, MOI.OBJECTIVE_LIMIT]
        fill!(Xsol, 0.0)
        fill!(lambda, 0.0)
        fill!(mult_x_U, 0.0)
        fill!(mult_x_L, 0.0)
    else
        @error "Unexpected status: $(status)"
    end

    return Xsol, lambda, mult_x_U, mult_x_L, p_slack
end


using LinearAlgebra

const maxval = 4
const numbits = 2
const data = [1, 1, 2, 3, 4]
const mhash = 4

""" holds all reports of clients. Each element is an array """
mutable struct Report{T}
   thedata
   encoded
   perturbed
   reconstructed  # unary encoding of reconstruction
   Report{T}() where {T} = new(data, missing, missing, missing)
   Report{T}(dataparam) where {T} = new(dataparam, missing, missing, missing)
end

encode(rs::Array{Report{T}}, epsilon) where {T} = [encode(r, epsilon) for r in rs]
perturb(rs::Array{Report{T}}, epsilon) where {T} = [perturb(r, epsilon) for r in rs]
reconstruct(rs::Array{Report{T}}, epsilon) where {T} = [reconstruct(r, epsilon) for r in rs]

function Base.show(io::IO, r::Report{T}) where {T}
    println(io, "Report{",T,"}")
    print(io, "Data: ")
    show(io, r.thedata)
    println(io)
    println(io, "Encoded: ")
    show(io, "text/plain", r.encoded)
    println(io)
    println(io, "Perturbed: ")
    show(io, "text/plain", r.perturbed)
    println(io)
    println(io, "Reconstructed: ")
    show(io, "text/plain", r.reconstructed)

end

#############################################################################
""" For reports that use unary encoding, RR on each bit"""
abstract type Unary end
UnaryReport = Report{Unary}

function encode(r::UnaryReport, epsilon)
   r.encoded = map(r.thedata) do d
      x = zeros(maxval) .- 1
      x[d] = 1
      x
   end
   r
end

function perturb(r::UnaryReport, epsilon)
   r.perturbed = map(r.encoded) do d
      [randomized_response(x, epsilon/2) for x in d]
   end
   r
end

function reconstruct(r::UnaryReport, epsilon)
    r.reconstructed = [(p .+ 1) ./ 2  for p in r.perturbed]
    r
end

#############################################################################
""" For reports that use the value representation, joint RR"""
abstract type RRUnary end
RRUnaryReport = Report{RRUnary}

function encode(r::RRUnaryReport, epsilon)
   r.encoded = r.thedata
   r
end

function perturb(r::RRUnaryReport, epsilon)
   r.perturbed = map(r.encoded) do d
      randomized_multi_response(d, maxval, epsilon)
   end
   r
end

function reconstruct(r::RRUnaryReport, epsilon)
   rrunary_matrx = [ exp(epsilon) 1 1 1;
      		    1  exp(epsilon) 1 1;
		    1 1  exp(epsilon) 1;
		    1 1 1  exp(epsilon)] ./ (maxval - 1 + exp(epsilon))
   r.reconstructed = map(r.perturbed) do d
       indicator = zeros(maxval)
       indicator[d] = 1
       rrunary_matrx \ indicator
   end
   r
end

#############################################################################
""" For reports that use binary encoding, RR on each bit
first bit is 1 if (d-1) & 1 == 1, second bit is 1 if (d-1) & 2 == 1
"""
abstract type Binary end
BinaryReport = Report{Binary}

function encode(r::BinaryReport, epsilon)
   r.encoded = map(r.thedata) do d
      x = zeros(numbits)
      x[1] = if (d == 2 || d == 4) 1 else -1 end
      x[2] = if (d == 3 || d == 4) 1 else -1 end
      x
   end
   r
end

function perturb(r::BinaryReport, epsilon)
   r.perturbed = map(r.encoded) do d
      [randomized_response(x, epsilon/2) for x in d]
   end
   r
end

function reconstruct(r::BinaryReport, epsilon)
   r.reconstructed = map(r.perturbed) do d
      indicator = zeros(maxval)
      pos = 1 + (d[1] > 0 ? 1 : 0) + 2 * (d[2] > 0 ? 1 : 0)
      indicator[pos] = 1
      p = exp(epsilon/2)/(1+exp(epsilon/2))
      thematrix = [p^2      p*(1-p)  p*(1-p)  (1-p)^2;
	           p*(1-p)  p^2      (1-p)^2  p*(1-p);
		   p*(1-p)  (1-p)^2  p^2      p*(1-p);
		   (1-p)^2  p*(1-p)  p*(1-p)  p^2]
      thematrix \ indicator
   end
   r
end

#############################################################################
""" For reports that choose a bit at random  Reports are tuples. First element is bit position, second is value"""
abstract type RBit end
RBitReport = Report{RBit}

rbit_matrix = [-1  1 -1 1;
	       -1 -1  1 1]

function encode(r::RBitReport, epsilon)
   r.encoded = map(r.thedata) do d
      chosen_bit = rand(1:size(rbit_matrix, 1))
      valuebit = rbit_matrix[chosen_bit, d]
      (chosen_bit, valuebit)
   end
   r
end

function perturb(r::RBitReport, epsilon)
   r.perturbed = map(r.encoded) do d
      pert = randomized_response(d[2], epsilon)
      (d[1], pert)
   end
   r
end

function reconstruct(r::RBitReport, epsilon)
   r.reconstructed = map(r.perturbed) do d
       rbit_matrix[d[1], :]  * d[2]
   end
   r
end

#############################################################################
""" For reports that use a random yes/no question. First element is the indicator
second is the answer """
abstract type YNBit end
YNBitReport = Report{YNBit}

function encode(r::YNBitReport, epsilon)
   yn_matrix = [rand([-1,1]) for _ in 1:mhash, _ in 1:maxval]
   r.encoded = map(r.thedata) do d
      chosen_bit = rand(1:size(yn_matrix, 1))
      valuebit = yn_matrix[chosen_bit, d]
      (copy(yn_matrix[chosen_bit,:]), valuebit)
   end
   r
end

function perturb(r::YNBitReport, epsilon)
   r.perturbed = map(r.encoded) do d
      pert = randomized_response(d[2], epsilon)
      (d[1], pert)
   end
   r
end

function reconstruct(r::YNBitReport, epsilon)
   r.reconstructed = map(r.perturbed) do d
       d[1] * d[2]
   end
   r
end

#############################################################################
""" For reports that use a balanced y/n questions where matrix of set indicators has 0 column sum
returns tuple where first element is row number and second is answer to that question
"""
abstract type Balanced end
BalancedReport = Report{Balanced}

balanced_matrix = [ 1  1  1   1;
	            1  1  1  -1;
	            1  1 -1   1;
	            1  1 -1  -1;
                    1 -1  1   1;
	            1 -1  1  -1;
	            1 -1 -1   1;
	            1 -1 -1  -1;
                   -1  1  1   1;
	           -1  1  1  -1;
	           -1  1 -1   1;
	           -1  1 -1  -1;
                   -1 -1  1   1;
	           -1 -1  1  -1;
	           -1 -1 -1   1;
	           -1 -1 -1  -1
	    ]

function encode(r::BalancedReport, epsilon)
   r.encoded = map(r.thedata) do d
      chosen_bit = rand(1:size(balanced_matrix, 1))
      valuebit = balanced_matrix[chosen_bit, d]
      (chosen_bit, valuebit)
   end
   r
end

function perturb(r::BalancedReport, epsilon)
   r.perturbed = map(r.encoded) do d
      pert = randomized_response(d[2], epsilon)
      (d[1], pert)
   end
   r
end

function reconstruct(r::BalancedReport, epsilon)
   r.reconstructed = map(r.perturbed) do d
       balanced_matrix[d[1], :]  * d[2]
   end
   r
end

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

""" applies differetially private randomized response to the +/- 1 bit. The output expected value is the bit"""
function randomized_response(thebit, epsilon)
   p = exp(epsilon)/(1 + exp(epsilon))
   multiplier = (exp(epsilon) + 1) / (exp(epsilon) - 1)
   multiplier * (if rand() <= p  1   else   -1   end) * thebit
end

function randomized_multi_response(theval, numchoices, epsilon)
   p = (exp(epsilon)-1)/(numchoices - 1 + exp(epsilon))
   result = if rand() <= p
      theval
   else
      rand(1:numchoices)
   end
   result
end

#########################################
# For testing
########################################

function test1()
   x = [randomized_response(-1, log(2)) for _ in 1:1000000]
   y = [randomized_response(1, log(2)) for _ in 1:1000000]
   println(count(s -> s==3, x)/length(x), " should be about 1/3")
   println(count(s -> s==3, y)/length(y), " should be about 2/3")
   println(sum(x)/length(x), " should be about -1")
   println(sum(y)/length(y), " should be about 1")
end

function test2()
   println("data is ", data)
   runary = UnaryReport()
   r_rrunary = RRUnaryReport()
   rbinary = BinaryReport()
   rbit = RBitReport()
   ryn = YNBitReport()
   rbal = BalancedReport()
   println("###########################")
   println("Unary")
   display(encode(runary, 1).encoded)
   println("Perturbed")
   display(perturb(runary, 1).perturbed)
   println("###########################")
   println("RRUnary")
   display(encode(r_rrunary, 1).encoded)
   println("Perturbed")
   display(perturb(r_rrunary, 1).perturbed)
   println("###########################")
   println("Binary")
   display(encode(rbinary, 1).encoded)
   println("Perturbed")
   display(perturb(rbinary, 1).perturbed)
   println("###########################")
   println("Rbit")
   display(encode(rbit, 1).encoded)
   println("Perturbed")
   display(perturb(rbit, 1).perturbed)
   println("###########################")
   println("YN")
   display(encode(ryn, 1).encoded)
   println("Perturbed")
   display(perturb(ryn, 1).perturbed)
   println("###########################")
   println("Balanced")
   display(encode(rbal, 1).encoded)
   println("Perturbed")
   display(perturb(rbal, 1).perturbed)
   println("###########################")
end

##################################################
#### Evaluation
##################################################

function evaluator(;epsilon=5, repetitions=100000, mydata=data)
    mymethods = [Unary, RRUnary, Binary, YNBit, Balanced]
    encoded_data = zeros(maxval)
    for d in mydata
       encoded_data[d] += 1
    end
    for m in mymethods
        allreports = [reconstruct(perturb(encode(Report{m}(mydata), epsilon), epsilon), epsilon) for _ in 1:repetitions]
	thecounts = [sum(r.reconstructed) for r in allreports]
	println("Method: ", m)
	println("Data")
	show(encoded_data)
	println()
	println("Mean")
	show(sum(thecounts)/length(thecounts))
	println()
	println("Variance")
	show(sum([(tc - encoded_data).^2 for tc in thecounts])/length(thecounts))
	println()
	println()
	println("----------")
	println()
    end
end


abstract type MultiplicativeStructure{Twisted,I} <: AbstractMatrix{I} end

_istwisted(::MultiplicativeStructure{T}) where {T} = T
_product(ms::MultiplicativeStructure, g, h) = _product(Val(_istwisted(ms)), g, h)

_product(::Val{false}, g, h) = g * h
_product(::Val{true}, g, h) = star(g) * h

struct ProductNotDefined <: Exception
    i::Any
    j::Any
    msg::Any
end

function Base.showerror(io::IO, ex::ProductNotDefined)
    print(io, "Product of elements $(ex.i) and $(ex.j) is not defined on the basis")
    print(io, " or the multiplicative structure could not be completed")
    if isdefined(ex, :msg)
        print(io, ": $(ex.msg)")
    end
    print(io, ".")
end

struct TrivialMStructure{Tw,I,B<:AbstractBasis} <: MultiplicativeStructure{Tw,I}
    basis::B
end

TrivialMStructure{Tw}(basis::AbstractBasis{T,I}) where {Tw,T,I} =
    TrivialMStructure{Tw,I,typeof(basis)}(basis)

basis(mstr::TrivialMStructure) = mstr.basis
Base.size(mstr::TrivialMStructure) = (l = length(basis(mstr)); (l, l))

Base.@propagate_inbounds function Base.getindex(
    mstr::TrivialMStructure,
    i::Integer,
    j::Integer,
)
    @boundscheck checkbounds(mstr, i, j)
    b = basis(mstr)
    g, h = b[i], b[j]
    gh = _product(mstr, g, h)
    gh in b || throw(ProductNotDefined(i, j, "$g · $h = $gh"))
    return b[gh]
end


#=
test_timeevolution:
- Julia version: 1.5.2
- Author: qcancam
- Date: 2021-01-03
=#

include("../src/MatrixProductStates.jl")

using .MatrixProductStates, SparseArrays, Arpack, PyCall, PyPlot

#np=pyimport("numpy")
#@pyimport matplotlib.pyplot as plt

g = 1.0; L = 3
T=100.0
N=1000.0
τ = T/N

function H_LZ(g, L)
    H = Array{Complex{Float64}, 2}[]
    id = [1  0;
          0  1]
    σˣ = [0  1;
          1  0]
    σᶻ = [1  0;
          0 -1]
    for i in 1:L-2
        push!(H,g*σᶻ⊗id+σˣ⊗id)
    end
    #=L-1=#
    push!(H,g*σᶻ⊗id+σˣ⊗id+g*id⊗σᶻ+id⊗σˣ)

    return H
end

function H_Z_MPO(L)
    id = [1  0;
          0  1]
    σˣ = [0  1;
          1  0]
    σᶻ = [1  0;
          0 -1]
    W_tnsr = zeros(Complex{Float64}, 2, 2, 2, 2)
    W_tnsr[1, 1, :, :] = id
    W_tnsr[2, 1, :, :] = σᶻ
    W_tnsr[2, 2, :, :] = id

    return MPO(W_tnsr, L)
end

#H = H_LZ(g, L)
ψ = randn(MPS{L, Complex{Float64}}, 100, 2)

glist=[-T/2+T/N*i for i in 1:N]
H1=[H_LZ(g, L)[1] for g in glist]
H2=[H_LZ(g, L)[2] for g in glist]
HLm1=[H_LZ(g, L)[L-1] for g in glist]

ϕ,time_list,expect_list= real_timedep_time_evolution(ψ, H1,H2,HLm1,H_Z_MPO(L), T,N,50)

#ϕ= real_time_evolution(ψ, H[1],H[2],H[L-1],H_TFIM_MPO(g,L), T,N,50)

fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
ax.plot(time_list, expect_list, "-", c="b", ms=10)
ax.set_xlabel("t")
ax.set_ylabel("⟨-∑σᶻ⟩")
#ax.set_title("2D Line Plot Sample")
#ax.grid(true)
plt.show()


module CompactBases

import Base: eltype, axes, size, ==, getindex, checkbounds, copyto!, similar, show, step
import Base.Broadcast: materialize

using ContinuumArrays
import ContinuumArrays: Basis, ℵ₁, Derivative, Inclusion, @simplify, simplifiable, _simplify,
    AbstractQuasiArrayApplyStyle

using QuasiArrays
import QuasiArrays: AbstractQuasiMatrix, QuasiAdjoint, MulQuasiArray,
    PInvQuasiMatrix, InvQuasiMatrix, QuasiDiagonal,
    BroadcastQuasiArray, SubQuasiArray, LazyQuasiMatrix

using IntervalSets

using LazyArrays

using LinearAlgebra
import LinearAlgebra: Matrix, dot
using SparseArrays
using FillArrays

using OffsetArrays
using BandedMatrices
using BlockBandedMatrices

using FastGaussQuadrature

using Formatting
using RecipesBase

vandermonde(B::Basis) = B[locs(B),:]

include("distributions.jl")
include("restricted_bases.jl")
include("types.jl")
include("unit_vectors.jl")
include("intervals.jl")
include("knot_sets.jl")
include("quadrature.jl")
include("skewtridiag.jl")

include("materialize_dsl.jl")

include("coulomb_derivative.jl")

include("fornberg.jl")
include("finite_differences.jl")
include("fd_operators.jl")
include("fd_derivatives.jl")

include("fedvr.jl")
include("fedvr_operators.jl")
include("fedvr_derivatives.jl")

include("bsplines.jl")
include("bspline_derivatives.jl")

include("inner_products.jl")
include("densities.jl")
include("linear_operators.jl")

"""
    centers(B)

Return the locations of the mass centers of all basis functions of
`B`; for orthogonal bases such as finite-differences and FE-DVR, this
is simply [`locs`](@ref), i.e. the location of the quadrature nodes.
"""
centers(B::BasisOrRestricted) = locs(B)

"""
    metric(B)

Returns the metric or mass matrix of the basis `B`, equivalent to
`S=B'B`.
"""
metric(B::BasisOrRestricted) = B'B

"""
    metric_shape(B)

Returns the shape of the metric or mass matrix of the basis `B`,
i.e. `Diagonal`, `BandedMatrix`, etc.
"""
metric_shape(B::BasisOrRestricted) = typeof(metric(B))

function matrix_element_metric(B::BasisOrRestricted)
    # For non-orthogonal bases, we need to apply the metric inverse,
    # after computing the action of a matrix representation of a
    # linear operator on a ket. However, when computing the inner
    # product with the bra afterwards, we again use the metric. Here,
    # we short-circuit this by combining them into them identity
    # operator I, if the operator metric equals the metric. For
    # orthogonal bases, with the integration weights not built into the
    # coefficients, and an identity matrix operator metric, we /do/
    # need to use the metric for the inner product.
    S = metric(B)
    oS = operator_metric(B)
    oS == S ? I : (oS \ S)
end

export AbstractFiniteDifferences,
    FiniteDifferences, StaggeredFiniteDifferences, ImplicitFiniteDifferences,
    Derivative, dot, QuasiDiagonal, Inclusion, .., distribution,
    FEDVR, Derivative, @elem, dot,
    BSpline,
    ArbitraryKnotSet, LinearKnotSet, ExpKnotSet,
    order, numintervals, numfunctions, nonempty_intervals,
    centers, vandermonde,
    metric, metric_shape, matrix_element_metric,
    FunctionProduct, Density,
    LinearOperator, DiagonalOperator, ShiftAndInvert

end # module


# 0 -> empty space, 1 -> O, 2 -> X
str = [' ', 'O', 'X']

function playTicTacToe(;ai::Int=1,method::String="minimax",abDepth::Int=4)
    # initialise board and display
    M = ttt_initialise()
    ttt_boardDisplay(M)

    # which player(s) are AI?
    if ai == 0
        isAI = [false false]
    elseif ai == 1
        isAI = rand()<0.5 ? [false true] : [true false]
    elseif ai == 2
        isAI = [true true]
    end

    # randomly pick starting player (sets symbols)
    player = rand((1,2))

    while true # take turns until the game ends
        # take a turn, either by human or AI
        if isAI[player]
            M = ttt_aiTurn!(M,player,method,abDepth)
        else
            M = ttt_takeTurn!(M,player)
        end

        ttt_boardDisplay(M) # show the board after the turn

        player = 3-player # swap player
        W = ttt_checkWinner(M) # check if someone has won

        if W!=0 # someone has won
            print("$(str[W+1])'s win\n")
            return nothing
        elseif ttt_gameOver(M) # board filled but no winner -> draw
            print("It's a draw\n")
            return nothing
        end
    end
end

function ttt_initialise()
    # empty board
    return zeros(Int8,3,3)
end

function ttt_boardDisplay(M)
    # Displays the board with symbols
    r1 = "         |   |   \n"
    r2(a) = "Row $a  $(str[M[a,1]+1]) | $(str[M[a,2]+1]) | $(str[M[a,3]+1]) \n"
    r3 = "      ___|___|___\n"
    print("Column 1   2   3\n",r1,r2(1),r3,
        r1,r2(2),r3,r1,r2(3),r1)
        print("Board eval: $(ttt_scoreHelper(M,1))\n")
end

function ttt_takeTurn!(M,player)
    # asks for player input and updates the game state
    coords = ttt_getMove(M)
    M[coords[1],coords[2]] = player
    return M
end

function ttt_getMove(M)
    # gets a move from a human player
    # TODO: More input validation

    while true # repeat until valid location selected
        print("Enter the row/column coordinates for your move separated by a space: ")
        # read a space-separated list of integers
        x = readdlm(IOBuffer(readline()),' ',Int64)
        if M[x[1],x[2]]==0 # must be an empty location
            return x
        end
    end
end

function ttt_aiTurn!(M,player,method,abDepth)
    # work out the best move and then plays it
    moves = 1:9
    if isequal(method,"minimaxAB")
        idx = getBestMoveAB(M,player,ttt_score,moves,ttt_movePossible,ttt_makeMove!,ttt_gameOver,abDepth)
    else
        idx = getBestMove(M,player,ttt_score,moves,ttt_movePossible,ttt_makeMove!,ttt_gameOver)
    end
    M[idx] = player
    return M
end

function ttt_score(M,player)
    win = ttt_checkWinner(M)
    if win!=0
        return win==player ? 100 : -100
    elseif ttt_gameOver(M)
        return 0
    else
        return ttt_scoreHelper(M,player)
    end

    # returns 0 if the game isn't over, 1 if the player wins,
    # -1 if the player loses
end

function ttt_scoreHelper(M,player)
    # See who has won the game

    # p1  and p2 are boolean arrays indicating where each player has played
    p = M.==player
    o = M.==3-player

    s = sum(filter(x -> x>=2 || x<=-2,
        [sum(p-o,dims=1) sum(p-o,dims=2)' sum(p[1:4:9]-o[1:4:9]) sum(p[3:2:7]-o[3:2:7])]))

    return s
end

function ttt_checkWinner(M)
    # See who has won the game

    # p1  and p2 are boolean arrays indicating where each player has played
    p1 = M.==1
    p2 = M.==2

    #   column of 3       or   row of 3         or main diagonal  or other diagonal
    if any(all(p1,dims=1)) | any(all(p1,dims=2)) | all(p1[1:4:9]) | all(p1[3:2:7])
        W = 1
    elseif any(all(p2,dims=1)) | any(all(p2,dims=2)) | all(p2[1:4:9]) | all(p2[3:2:7])
        W = 2
    else
        W = 0 # no winner or draw
    end

    return W
end

function ttt_makeMove!(N,i,player)
    N[i] = player
    return N
end

function ttt_movePossible(M,i)
    return M[i]==0
end

function ttt_gameOver(M)
    # if all the squares are taken (filled with 1s/2s) the game ends
    return all(M.>0) || ttt_checkWinner(M)!=0
end


using LinearAlgebra
@testset "degopt_poly" begin
    x1a = [3; 4.0]
    x1b = [0.1; 0.3]
    x2a = [0.1; 0.1; -0.1]
    x2b = [0.3; 0; -0.1]
    x = [(x1a, x1b), (x2a, x2b)]

    z = [1.0; -0.1; 0.1; 0.01]
    (graph, crefs) = graph_degopt(x, z)

    x0 = 3.1415
    ev1 = eval_graph(graph, x0, comporder = get_topo_order_degopt(2))

    # Manual evaluation
    β1 = (x1a[1] + x1a[2] * x0) * (x1b[1] + x1b[2] * x0)
    β2 =
        (x2a[1] + x2a[2] * x0 + x2a[3] * β1) *
        (x2b[1] + x2b[2] * x0 + x2b[3] * β1)
    ev2 = z[1] + z[2] * x0 + z[3] * β1 + z[4] * β2
    @test ev1 ≈ ev2

    # Also check that the crefs are correct

    # Modify the graph: set almost all coeffs to one
    T = Float64
    t = ones(T, size(crefs, 1))
    t[2] = 2.0
    set_coeffs!(graph, t, crefs)

    # Create a new graph object graph2: with coeffs almost all zero
    x1a = [1; 2.0]
    x1b = [1; 1.0]
    x2a = [1.0; 1.0; 1.0]
    x2b = [1.0; 1.0; 1.0]
    z = ones(T, 4)
    x = [(x1a, x1b), (x2a, x2b)]
    (graph2, crefs2) = graph_degopt(x, z)

    x1 = 0.3
    @test eval_graph(graph, x1, comporder = get_topo_order_degopt(2)) ==
          eval_graph(graph2, x1)

    for k = 1:7
        x = Vector{Tuple{Vector{T},Vector{T}}}(undef, k)
        for i = 1:k
            x[i] = (collect(1.0:i+1), collect(-1.0:-1:-(i + 1)))
        end
        z = collect(range(4, 5, length = k + 2))
        (graph, cref) = graph_degopt(x, z)

        (xx, zz) = get_degopt_crefs(k)
        for i = 1:k
            for N = 1:2
                @test get_coeffs(graph, xx[i][N]) == x[i][N]
            end
        end
        @test get_coeffs(graph, zz) == z
    end

    # Check the inteface with all ones
    (graph, crefs) = graph_degopt(3)
    x = randn(3)
    a = (1 .+ x) .^ 2
    b = (1 .+ x .+ a) .^ 2
    c = (1 .+ x .+ a .+ b) .^ 2
    g = 1 .+ x .+ a .+ b .+ c
    @test eval_graph(graph, x) ≈ g
end


using Ferrite
using BlockArrays, SparseArrays, LinearAlgebra

function create_cook_grid(nx, ny)
    corners = [Vec{2}((0.0,   0.0)),
               Vec{2}((48.0, 44.0)),
               Vec{2}((48.0, 60.0)),
               Vec{2}((0.0,  44.0))]
    grid = generate_grid(Triangle, (nx, ny), corners);
    # facesets for boundary conditions
    addfaceset!(grid, "clamped", x -> norm(x[1]) ≈ 0.0);
    addfaceset!(grid, "traction", x -> norm(x[1]) ≈ 48.0);
    return grid
end;

function create_values(interpolation_u, interpolation_p)
    # quadrature rules
    qr      = QuadratureRule{2,RefTetrahedron}(3)
    face_qr = QuadratureRule{1,RefTetrahedron}(3)

    # geometric interpolation
    interpolation_geom = Lagrange{2,RefTetrahedron,1}()

    # cell and facevalues for u
    cellvalues_u = CellVectorValues(qr, interpolation_u, interpolation_geom)
    facevalues_u = FaceVectorValues(face_qr, interpolation_u, interpolation_geom)

    # cellvalues for p
    cellvalues_p = CellScalarValues(qr, interpolation_p, interpolation_geom)

    return cellvalues_u, cellvalues_p, facevalues_u
end;

function create_dofhandler(grid, ipu, ipp)
    dh = DofHandler(grid)
    push!(dh, :u, 2, ipu) # displacement
    push!(dh, :p, 1, ipp) # pressure
    close!(dh)
    return dh
end;

function create_bc(dh)
    dbc = ConstraintHandler(dh)
    add!(dbc, Dirichlet(:u, getfaceset(dh.grid, "clamped"), (x,t) -> zero(Vec{2}), [1,2]))
    close!(dbc)
    t = 0.0
    update!(dbc, t)
    return dbc
end;

struct LinearElasticity{T}
    G::T
    K::T
end

function doassemble(cellvalues_u::CellVectorValues{dim}, cellvalues_p::CellScalarValues{dim},
                    facevalues_u::FaceVectorValues{dim}, K::SparseMatrixCSC, grid::Grid,
                    dh::DofHandler, mp::LinearElasticity) where {dim}

    f = zeros(ndofs(dh))
    assembler = start_assemble(K, f)
    nu = getnbasefunctions(cellvalues_u)
    np = getnbasefunctions(cellvalues_p)

    fe = PseudoBlockArray(zeros(nu + np), [nu, np]) # local force vector
    ke = PseudoBlockArray(zeros(nu + np, nu + np), [nu, np], [nu, np]) # local stiffness matrix

    # traction vector
    t = Vec{2}((0.0, 1/16))
    # cache ɛdev outside the element routine to avoid some unnecessary allocations
    ɛdev = [zero(SymmetricTensor{2, dim}) for i in 1:getnbasefunctions(cellvalues_u)]

    for cell in CellIterator(dh)
        fill!(ke, 0)
        fill!(fe, 0)
        assemble_up!(ke, fe, cell, cellvalues_u, cellvalues_p, facevalues_u, grid, mp, ɛdev, t)
        assemble!(assembler, celldofs(cell), fe, ke)
    end

    return K, f
end;

function assemble_up!(Ke, fe, cell, cellvalues_u, cellvalues_p, facevalues_u, grid, mp, ɛdev, t)

    n_basefuncs_u = getnbasefunctions(cellvalues_u)
    n_basefuncs_p = getnbasefunctions(cellvalues_p)
    u▄, p▄ = 1, 2
    reinit!(cellvalues_u, cell)
    reinit!(cellvalues_p, cell)

    # We only assemble lower half triangle of the stiffness matrix and then symmetrize it.
    @inbounds for q_point in 1:getnquadpoints(cellvalues_u)
        for i in 1:n_basefuncs_u
            ɛdev[i] = dev(symmetric(shape_gradient(cellvalues_u, q_point, i)))
        end
        dΩ = getdetJdV(cellvalues_u, q_point)
        for i in 1:n_basefuncs_u
            divδu = shape_divergence(cellvalues_u, q_point, i)
            δu = shape_value(cellvalues_u, q_point, i)
            for j in 1:i
                Ke[BlockIndex((u▄, u▄), (i, j))] += 2 * mp.G * ɛdev[i] ⊡ ɛdev[j] * dΩ
            end
        end

        for i in 1:n_basefuncs_p
            δp = shape_value(cellvalues_p, q_point, i)
            for j in 1:n_basefuncs_u
                divδu = shape_divergence(cellvalues_u, q_point, j)
                Ke[BlockIndex((p▄, u▄), (i, j))] += -δp * divδu * dΩ
            end
            for j in 1:i
                p = shape_value(cellvalues_p, q_point, j)
                Ke[BlockIndex((p▄, p▄), (i, j))] += - 1/mp.K * δp * p * dΩ
            end

        end
    end

    symmetrize_lower!(Ke)

    # We integrate the Neumann boundary using the facevalues.
    # We loop over all the faces in the cell, then check if the face
    # is in our `"traction"` faceset.
    @inbounds for face in 1:nfaces(cell)
        if onboundary(cell, face) && (cellid(cell), face) ∈ getfaceset(grid, "traction")
            reinit!(facevalues_u, cell, face)
            for q_point in 1:getnquadpoints(facevalues_u)
                dΓ = getdetJdV(facevalues_u, q_point)
                for i in 1:n_basefuncs_u
                    δu = shape_value(facevalues_u, q_point, i)
                    fe[i] += (δu ⋅ t) * dΓ
                end
            end
        end
    end
end

function symmetrize_lower!(K)
    for i in 1:size(K,1)
        for j in i+1:size(K,1)
            K[i,j] = K[j,i]
        end
    end
end;

function solve(ν, interpolation_u, interpolation_p)
    # material
    Emod = 1.
    Gmod = Emod / 2(1 + ν)
    Kmod = Emod * ν / ((1+ν) * (1-2ν))
    mp = LinearElasticity(Gmod, Kmod)

    # grid, dofhandler, boundary condition
    n = 50
    grid = create_cook_grid(n, n)
    dh = create_dofhandler(grid, interpolation_u, interpolation_p)
    dbc = create_bc(dh)

    # cellvalues
    cellvalues_u, cellvalues_p, facevalues_u = create_values(interpolation_u, interpolation_p)

    # assembly and solve
    K = create_sparsity_pattern(dh);
    K, f = doassemble(cellvalues_u, cellvalues_p, facevalues_u, K, grid, dh, mp);
    apply!(K, f, dbc)
    u = Symmetric(K) \ f;

    # export
    filename = "cook_" * (isa(interpolation_u, Lagrange{2,RefTetrahedron,1}) ? "linear" : "quadratic") *
                         "_linear"
    vtk_grid(filename, dh) do vtkfile
        vtk_point_data(vtkfile, dh, u)
    end
    return u,dh
end

linear    = Lagrange{2,RefTetrahedron,1}()
quadratic = Lagrange{2,RefTetrahedron,2}()

u,dh = solve(0.4999999, quadratic, linear);


function decompose_CNot_41()
    # Reference: https://doi.org/10.1109/DSD.2018.00005

    return Dict{String, Any}(

    "num_qubits" => 4,
    "maximum_depth" => 10,
    "elementary_gates" => ["H_1", "H_2", "H_3", "CNot_1_3", "CNot_4_3", "Identity"],
    "target_gate" => QCOpt.get_full_sized_gate("CNot_4_1", 4),
    "objective" => "minimize_depth",
    )

end

function decompose_quantum_fulladder()
    #Reference-1: https://doi.org/10.1109/DATE.2005.249
    #Reference-2: https://doi.org/10.1109/TCAD.2005.858352

    num_qubits = 4

    function target_gate_1()

        CV_1_2 = QCOpt.get_full_sized_gate("CV_1_2", num_qubits);
        CV_4_2 = QCOpt.get_full_sized_gate("CV_4_2", num_qubits);
        CVdagger_1_2 = QCOpt.get_full_sized_gate("CVdagger_1_2", num_qubits);
        CVdagger_3_2 = QCOpt.get_full_sized_gate("CVdagger_3_2", num_qubits);
        CNot_3_1 = QCOpt.get_full_sized_gate("CNot_3_1", num_qubits);
        CNot_4_3 = QCOpt.get_full_sized_gate("CNot_4_3", num_qubits);
        CNot_2_4 = QCOpt.get_full_sized_gate("CNot_2_4", num_qubits);

        return CNot_2_4 * CV_4_2 * CVdagger_1_2 * CVdagger_3_2 * CNot_4_3 * CNot_3_1 * CV_1_2
    end

    return Dict{String, Any}(

    "num_qubits" => num_qubits,
    "maximum_depth" => 7,
    "elementary_gates" => ["CV_1_2", "CV_4_2", "CV_3_2", "CVdagger_1_2", "CVdagger_4_2", "CVdagger_3_2", "CNot_3_1", "CNot_4_3", "CNot_2_4", "CNot_4_1", "Identity"],
    "target_gate" => target_gate_1(),
    "objective" => "minimize_depth",
    )

end

function decompose_double_toffoli()

    # Reference: https://doi.org/10.1109/TCAD.2005.858352

    num_qubits = 4

    function target_gate()
        CV_3_4 = QCOpt.get_full_sized_gate("CV_3_4", num_qubits);
        CV_2_4 = QCOpt.get_full_sized_gate("CV_2_4", num_qubits);
        CVdagger_2_4 = QCOpt.get_full_sized_gate("CVdagger_2_4", num_qubits);
        CNot_1_3 = QCOpt.get_full_sized_gate("CNot_1_3", num_qubits);
        CNot_3_2 = QCOpt.get_full_sized_gate("CNot_3_2", num_qubits);

        return CV_2_4 * CNot_1_3 * CNot_3_2 * CV_3_4 * CVdagger_2_4 * CNot_3_2 * CNot_1_3
    end

    return Dict{String, Any}(

    "num_qubits" => 4,
    "maximum_depth" => 7,
    "elementary_gates" => ["CV_1_2", "CV_2_4", "CV_3_4", "CVdagger_1_2", "CVdagger_2_4", "CVdagger_3_4", "CNot_1_3", "CNot_3_2", "CNot_2_3", "Identity"],
    "target_gate" => target_gate(),
    "objective" => "minimize_depth"
    )

end

function decompose_double_peres()

    # Reference: https://doi.org/10.1109/TCAD.2005.858352

    num_qubits = 4

    function target_gate()
        CV_1_4 = QCOpt.get_full_sized_gate("CV_1_4", num_qubits);
        CV_2_4 = QCOpt.get_full_sized_gate("CV_2_4", num_qubits);
        CV_3_4 = QCOpt.get_full_sized_gate("CV_3_4", num_qubits);
        CVdagger_3_4 = QCOpt.get_full_sized_gate("CVdagger_3_4", num_qubits);
        CNot_1_2 = QCOpt.get_full_sized_gate("CNot_1_2", num_qubits);
        CNot_2_3 = QCOpt.get_full_sized_gate("CNot_2_3", num_qubits);

        return CV_2_4 * CNot_1_2 * CV_3_4 * CV_1_4 * CNot_2_3 * CVdagger_3_4
    end

    return Dict{String, Any}(

    "num_qubits" => 4,
    "maximum_depth" => 7,
    "elementary_gates" => ["CV_1_4", "CV_2_4", "CV_3_4", "CVdagger_1_4", "CVdagger_2_4", "CVdagger_3_4", "CNot_1_2", "CNot_2_3", "Identity"],
    "target_gate" => target_gate(),
    "objective" => "minimize_depth"
    )

end


using FractionalDelayFilter
using Test

@testset "Max filter order" begin
    delaystep = 0.25
    delays = 0.0:delaystep:5.0
    numdelays = length(delays)
    fdfcoefs = Vector{Vector{Float64}}(undef, numdelays)
    nonzeroindices = Vector{Int}(undef, numdelays)
    filterorders = Vector{Int}(undef, numdelays)
    sinewave = sin.(0:π/40:4π)

    maxorder = 3
    @testset "Odd value: $maxorder" begin
        for delayindex = 1:numdelays
            filterorders[delayindex] = filtord(delays[delayindex], maxorder=maxorder)
            (fdfcoefs[delayindex], nonzeroindices[delayindex]) = getfdfcoef(filterorders[delayindex], delays[delayindex])
        end

        delayedwave = zeros(Float64, length(sinewave)*2, numdelays)
        for delayindex = 1:numdelays
            delayedwave[1+nonzeroindices[delayindex]:length(sinewave)+nonzeroindices[delayindex]+length(fdfcoefs[delayindex])-2,delayindex] = fdfilter(sinewave, fdfcoefs[delayindex])
        end
        @test sortperm(delayedwave[7,:], rev=true) == collect(1:numdelays)
    end

    maxorder = 4
    @testset "Even value: $maxorder" begin
        for delayindex = 1:numdelays
            filterorders[delayindex] = filtord(delays[delayindex], maxorder=maxorder)
            (fdfcoefs[delayindex], nonzeroindices[delayindex]) = getfdfcoef(filterorders[delayindex], delays[delayindex])
        end

        delayedwave = zeros(Float64, length(sinewave)*2, numdelays)
        for delayindex = 1:numdelays
            delayedwave[1+nonzeroindices[delayindex]:length(sinewave)+nonzeroindices[delayindex]+length(fdfcoefs[delayindex])-2,delayindex] = fdfilter(sinewave, fdfcoefs[delayindex])
        end
        @test sortperm(delayedwave[7,:], rev=true) == collect(1:numdelays)
    end

end


#### SVD, i.e. UDV decomposition
decompose_udv!(A::AbstractMatrix{<:Number}) = LinearAlgebra.LAPACK.gesvd!('A','A',A) # was gesdd! at some point
# decompose_udv!(A::AbstractMatrix{<:Number}) = LinearAlgebra.LAPACK.gesdd!('A',A) # was gesdd! at some point
# using JacobiSVD # experimental package
# decompose_udv!(A::AbstractMatrix{<:Number}) = JacobiSVD.gesvj!('G','U','V', A) # was gesdd! at some point
decompose_udv(A::AbstractMatrix{T}) where T<:Number = decompose_udv!(copy(A))

#### QR, i.e. UDT decomposition
function decompose_udt(A::AbstractMatrix{C}) where C<:Number
  F = qr(A, Val(true))
  p = F.p
  @views p[p] = collect(1:length(p))
  # D = abs.(real(diag(triu(R))))
  D = abs.(real(diag(F.R)))
  T = (Diagonal(1 ./ D) * F.R)[:, p]
  return Matrix(F.Q), D, T
end

function decompose_udt!(A::AbstractMatrix{C}, D) where C<:Number
  n = length(D)
  F = qr!(A, Val(true))
  R = F.R # F.R is of regular matrix type

  @views F.p[F.p] = 1:n

  @inbounds for i in 1:n
    D[i] = abs(real(R[i,i]))
  end

  # This (very!) strangely increases the runtime of measure_tdgfs! by a ton....
  # @inbounds for i in diagind(R)
  #   R[i] = sign(real(R[i]))
  # end

  lmul!(Diagonal(1 ./ D), R)

  return Matrix(F.Q), R[:, F.p] # Q, (D is modified in-place), T  # was full(F.Q) before upgrade
end

#### Other
# See https://discourse.julialang.org/t/asymmetric-speed-of-in-place-sparse-dense-matrix-product/10256/3
import LinearAlgebra.mul!
function mul!(C::StridedMatrix, X::StridedMatrix, A::SparseMatrixCSC)
    mX, nX = size(X)
    nX == A.m || throw(DimensionMismatch())
    fill!(C, zero(eltype(C)))
    rowval = A.rowval
    nzval = A.nzval
    @inbounds for  col = 1:A.n, k=A.colptr[col]:(A.colptr[col+1]-1)
        ki=rowval[k]
        kv=nzval[k]
        for multivec_row=1:mX
            C[multivec_row, col] += X[multivec_row, ki] * kv
        end
    end
    C
end

myrdiv!(dest, a, b) = copyto!(dest, adjoint(adjoint(b) \ adjoint(a)))

LinearAlgebra.Diagonal(I::UniformScaling{T}) where T <: Number = I.λ

function eye!(x)
    @assert ndims(x) == 2
    @assert size(x,1) == size(x,2)

    o = one(eltype(x))
    z = zero(eltype(x))

    for j in 1:size(x,2)
        for i in 1:size(x,1)
            @inbounds x[i, j] = i == j ? o : z
        end
    end
    x
end

"""
Kronecker Delta
"""
@inline kd(a,b) = a == b

##############################################################
#
#                   QR / UDT
#
##############################################################
"""
Safely multiply two UDT decompositions
"""
function multiply_safely(Ul,Dl,Tl, Ur,Dr,Tr)
  mat = Tl * Ur
  lmul!(Diagonal(Dl), mat)
  rmul!(mat, Diagonal(Dr))
  U, D, T = decompose_udt(mat)
  return Ul*U, D, T*Tr
end

"""

  inv_udt(U, D, T) -> result

Calculate (UDT)^(-1).
"""
function inv_udt(U,D,T)
  res = similar(U)
  inv_udt!(res, U, D, T)
  res
end

"""

  inv_udt!(res, U, D, T) -> nothing

Calculate (UDT)^(-1) and store result in `res`
"""
function inv_udt!(res, U,D,T)
  tmp = similar(U)
  ldiv!(tmp, lu(T), Diagonal(1 ./ D))
  mul!(res, tmp, U')
  nothing
end

"""

  inv_one_plus_udt(U, D, T) -> result

Stable calculation of [1 + UDT]^(-1):

  * Use one intermediate UDT decomposition.

Faster but less accurate than the scalettar approach.

Consider `inv_one_plus_udt!` as an efficient (not one-to-one) replacement.
"""
function inv_one_plus_udt(U,D,T)
  @warn "Calling somewhat inefficient and potentially inaccurate `inv_one_plus_udt`"

  m = U' / T
  m[diagind(m)] .+= D
  u,d,t = decompose_udt(m)
  u = U * u
  t = t * T
  ldiv!(m, lu!(t), Diagonal(1 ./ d))
  m * u'
end

"""

  inv_one_plus_udt!(mc, res, U, D, T) -> nothing

Stable calculation of [1 + UDT]^(-1):

  * Use one intermediate UDT decomposition.

Uses preallocated memory in `mc`. Writes the result into `res`.

Much faster (~50%) than `inv_one_plus_udt_scalettar!` but less accurate.
"""
function inv_one_plus_udt!(mc, res, U,D,T)
  @warn "Calling potentially inaccurate `inv_one_plus_udt!`"

  d = mc.s.d
  u = mc.s.tmp
  t = mc.s.tmp2

  m = U' / T
  m[diagind(m)] .+= D

  utmp,ttmp = decompose_udt!(m, d)
  mul!(u, U, utmp)
  mul!(t, ttmp, T)

  ldiv!(m, lu!(t), Diagonal(1 ./ d))

  mul!(res, m, u')
  nothing
end

"""

  inv_one_plus_two_udts!(mc, U, D, T, Ul, Dl, Tl, Ur, Dr, Tr) -> nothing

Stable calculation of [1 + UlDlTl(UrDrTr)^†]^(-1).

Uses preallocated memory in `mc`. Writes the result into `U`, `D`, and `T`.
"""
function inv_one_plus_two_udts!(mc, U,D,T, Ul,Dl,Tl, Ur,Dr,Tr)
  s = mc.s
  tmp = mc.s.tmp
  tmp2 = mc.s.tmp2
  tmp3 = mc.s.curr_U

  mul!(tmp, Tl, adjoint(Tr))
  rmul!(tmp, Diagonal(Dr))
  lmul!(Diagonal(Dl), tmp)
  U1, T1 = decompose_udt!(tmp, s.D)

  mul!(tmp3, Ul, U1)
  mul!(tmp2, T1, adjoint(Ur))
  mul!(tmp, adjoint(tmp3), inv(tmp2))

  tmp .+= Diagonal(s.D)

  u, t = decompose_udt!(tmp, D)

  mul!(tmp, t, tmp2)
  copyto!(U, inv(tmp))

  mul!(tmp, tmp3, u)
  copyto!(T, adjoint(tmp))

  D .= 1 ./ D

  nothing
end

"""

  inv_one_plus_two_udts!(mc, res, Ul, Dl, Tl, Ur, Dr, Tr) -> nothing

Stable calculation of [1 + UlDlTl(UrDrTr)^†]^(-1).

Uses preallocated memory in `mc`. Writes the result into `res`.
"""
function inv_one_plus_two_udts!(mc, res, Ul,Dl,Tl, Ur,Dr,Tr)
  inv_one_plus_two_udts!(mc, mc.s.U, mc.s.d, mc.s.T, Ul, Dl, Tl, Ur, Dr, Tr)
  rmul!(mc.s.U, Diagonal(mc.s.d))
  mul!(res, mc.s.U, mc.s.T)
  nothing
end

"""

  inv_one_plus_udt_scalettar!(mc, res, Ua, Da, Ta, Ub, Db, Tb) -> nothing

Stable calculation of [1 + UDT]^(-1):

  * Separate scales larger and smaller than unity
  * Use two intermediate UDT decompositions.

Uses preallocated memory in `mc`. Writes the result into `res`.
"""
function inv_one_plus_udt_scalettar!(mc, res, U,D,T)
  d = mc.s.d
  r = mc.s.tmp
  l = mc.s.tmp2

  Dp = max.(D,1.)
  Dm = min.(D,1.)

  Dp .\= 1
  Dpinv = Dp # renaming

  ldiv!(l, lu(T), Diagonal(Dpinv)) # Don't use lu!(T) because T is input

  mul!(r, U, Diagonal(Dm))
  r .+= l

  u, t = decompose_udt!(r, d)

  ldiv!(r, lu!(t), Diagonal(1 ./ d))
  mul!(l, r, u')

  lmul!(Diagonal(Dpinv), l)
  u, t = decompose_udt!(l, d)

  ldiv!(l, lu(T), u)

  rmul!(l, Diagonal(d))
  mul!(res, l, t)
  nothing
end

"""

  inv_sum_udts(Ua, Da, Ta, Ub, Db, Tb) -> U, D, T

Stable calculation of [UaDaTa + UbDbTb]^(-1):

  * Use one intermediate UDT decompositions.

Faster but less accurate than `inv_sum_udts_scalettar`.

Consider `inv_sum_udts!` as an efficient (not one-to-one) replacement.
"""
function inv_sum_udts(Ua,Da,Ta,Ub,Db,Tb)
  @warn "Calling somewhat inefficient and potentially inaccurate `inv_sum_udts`"

  m1 = Ta / Tb
  lmul!(Diagonal(Da), m1)

  m2 = Ua' * Ub
  rmul!(m2, Diagonal(Db))

  u,d,t = decompose_udt(m1 + m2)

  mul!(m1, Ua, u)
  mul!(m2, t, Tb)

  return inv(m2), 1 ./ d, m1'
end

"""

  inv_sum_udts!(mc, res, Ua, Da, Ta, Ub, Db, Tb) -> nothing

Stable calculation of [UaDaTa + UbDbTb]^(-1):

  * Use one intermediate UDT decompositions.

Uses preallocated memory in `mc`. Writes the result into `res`.

Much faster (~40%) than `inv_sum_udts_scalettar!` but less accurate.
"""
function inv_sum_udts!(mc, res, Ua,Da,Ta,Ub,Db,Tb)
  @warn "Calling potentially inaccurate `inv_sum_udts!`"

  d = mc.s.d
  m1 = mc.s.tmp
  m2 = mc.s.tmp2
  tmp = mc.s.U

  m1 = Ta / Tb
  lmul!(Diagonal(Da), m1)

  mul!(m2, Ua', Ub)
  rmul!(m2, Diagonal(Db))

  u,t = decompose_udt!(m1 + m2, d)

  mul!(m1, Ua, u)
  mul!(m2, t, Tb)

  ldiv!(tmp, lu!(m2), Diagonal(1 ./ d))
  mul!(res, tmp, m1')

  nothing
end

"""

  inv_sum_udts_scalettar(Ua, Da, Ta, Ub, Db, Tb) -> U, D, T

Stable calculation of [UaDaTa + UbDbTb]^(-1):

  * Separate scales larger and smaller than unity
  * Use two intermediate UDT decompositions.

Consider `inv_sum_udts_scalettar!` as an efficient (not one-to-one) replacement.
"""
function inv_sum_udts_scalettar(Ua, Da, Ta, Ub, Db, Tb)
    @warn "Calling somewhat inefficient `inv_sum_udts_scalettar`"

    d=length(Da)

    # separating scales larger and smaller than unity
    Dap = max.(Da,1.)
    Dam = min.(Da,1.)
    Dbp = max.(Db,1.)
    Dbm = min.(Db,1.)

    # mat1 = Dam * Vda * Vdb' / Dbp
    mat1 = Ta * inv(Tb)
    @inbounds for j in 1:d, k in 1:d
        mat1[j,k]=mat1[j,k] * Dam[j]/Dbp[k]
    end

    # mat2 = 1/(Dap) * Ua' * Ub * Dbm
    mat2 = Ua' * Ub
    @inbounds for j in 1:d, k in 1:d
        mat2[j,k]=mat2[j,k] * Dbm[k]/Dap[j]
    end

    #mat1 = mat1 + mat2
    mat1 = mat1 + mat2

    # decompose mat1: U, D, T
    U, D, T = decompose_udt(mat1)

    # invert inner part: mat1 = (U D T)^(-1) = mat1^(-1)
    # was UDT_to_mat!(mat1, U, D, T, invert=true)
    lmul!(Diagonal(D), T)
    ldiv!(mat1, lu!(T), U') # mat1 = T \ (U')

    # mat1 = 1/Dbp * mat1 /Dap
    @inbounds for j in 1:d, k in 1:d
        mat1[j,k]=mat1[j,k] / Dbp[j] / Dap[k]
    end

    #mat1 = U D T
    U, D, T = decompose_udt(mat1)

    # U = Tb^(-1) * U , T = T * Ua'
    mul!(mat1, inv(Tb), U)
    mul!(mat2, T, Ua')
    U = mat1
    T = mat2

    return U, D, T
end

"""

  inv_sum_udts_scalettar!(mc, res, Ua, Da, Ta, Ub, Db, Tb) -> nothing

Stable calculation of [UaDaTa + UbDbTb]^(-1):

  * Separate scales larger and smaller than unity
  * Use two intermediate UDT decompositions.

Uses preallocated memory in `mc`. Writes the result into `res`.
"""
function inv_sum_udts_scalettar!(mc, res, Ua, Da, Ta, Ub, Db, Tb)
    # optimization rounds: 1
    mat1 = mc.s.tmp
    mat2 = mc.s.U
    D = mc.s.d
    to = mc.a.to

    d=length(Da)

    # separating scales larger and smaller than unity
    Dap = max.(Da,1.)
    Dam = min.(Da,1.)
    Dbp = max.(Db,1.)
    Dbm = min.(Db,1.)

    # mat1 = Dam * Vda * Vdb' / Dbp
    mat1 = Ta / Tb
    @inbounds for j in 1:d, k in 1:d
        mat1[j,k]=mat1[j,k] * Dam[j]/Dbp[k]
    end

    # mat2 = 1/(Dap) * Ua' * Ub * Dbm
    mul!(mat2, Ua', Ub)
    @inbounds for j in 1:d, k in 1:d
        mat2[j,k]=mat2[j,k] * Dbm[k]/Dap[j]
    end

    # mat1 = mat1 + mat2
    mat1 .+= mat2

    # decompose mat1: U, D, T
    U, T = decompose_udt!(mat1, D)

    # invert and combine inner part: mat1 = (U D T)^(-1)
    lmul!(Diagonal(D), T)
    ldiv!(mat1, lu!(T), U') # mat1 = T \ (U')

    # mat1 = 1/Dbp * mat1 /Dap
    @inbounds for j in 1:d, k in 1:d
        mat1[j,k]=mat1[j,k] / Dbp[j] / Dap[k]
    end

    #mat1 = U D T
    U, T = decompose_udt!(mat1, D)

    # U = Tb^(-1) * U , T = T * Ua'
    ldiv!(mat1, lu!(Tb), U) # mat1 = Tb \ U
    mul!(mat2, T, Ua')
    U = mat1
    T = mat2

    # combine UDT into res
    rmul!(U, Diagonal(D))
    mul!(res, U, T)
    nothing
end

"""
  UDT_to_mat!(mat, U, D, T[; invert=false]) -> nothing

Combine given UDT to a matrix. Result will be stored in `mat`.

If `invert = true` (UDT)^(-1) is calculated.

Note that if parts of the input UDT may be overwritten it's
more efficient to inline this function and remove unnecessary copies!
"""
function UDT_to_mat!(mat, U, D, T; invert=false)
    # TODO: Check where this is called and inline manually for better speed. (remove copy())
    @warn "UDT_to_mat! probably shouldn't be called here" # TODO remove once checked
    if !invert
        mat1 = copy(U)
        rmul!(mat1, Diagonal(D))
        mul!(mat,mat1,T)
    else # (DT)^-1 * U^dagger
        mat1 = copy(T)
        lmul!(Diagonal(D), mat1)
        mat .= mat1 \ U'
    end
    nothing
end

function UDT_to_mat(U, D, T; kw...)
  res = copy(U)
  UDT_to_mat!(res, U, D, T; kw...)
  res
end

##############################################################
#
#                   SVD / UDV(t)
#
##############################################################
# multiplies two UDVds -> UDVd
function multiply_safely_udv(Ul,Dl,Vdl,Ur,Dr,Vdr)
  tmp = adjoint(Vdl) * Ur
  rmul!(tmp, Diagonal(Dr))
  lmul!(Diagonal(Dl), tmp)
  U, D, Vd = decompose_udv!(tmp)
  U = Ul * U
  Vd = Vd * Vdr
  return U, D, Vd
end

# Calculates (UDVd)^-1, where U, D, Vd come from SVD decomp.
function inv_udv(U,D,Vd)
  m = copy(Vd')
  rmul!(m, Diagonal(1 ./ D))
  res = similar(m)
  mul!(res, m, U')
  res
end

# Calculates (UDVd)^-1, where U, D, Vd come from SVD decomp.
function inv_udv!(res, U,D,Vd)

  # copy here isn't necessary but Vd would be overwritten if we drop it
  # also, m could be preallocated in stack if necessary
  m = copy(Vd')
  rmul!(m, Diagonal(1 ./ D))
  mul!(res, m, U')
  nothing
end

# Calculates (1 + UDVd)^-1, where U, D, Vd come from SVD decomp.
# !! Breaks down for large spread in D (i.e. low temperatures).
function inv_one_plus_udv(U,D,Vd)
  inner = copy(Vd')
  inner .+= U * Diagonal(D)
  I = decompose_udv!(inner)
  u = copy(adjoint(I[3] * Vd))
  d = 1 ./ I[2]
  vd = adjoint(I[1])

  rmul!(u,Diagonal(d))
  u * vd
end

# same as inv_one_plus_udt but separating both U AND Vd from D
# !! Breaks down for large spread in D (i.e. low temperatures). Slightly better than normal version.
function inv_one_plus_udv_alt(U,D,Vd)
  inner = copy((Vd*U)')
  inner[diagind(inner)] .+= D
  u, d, vd = decompose_udv!(inner)

  t1 = adjoint(vd*Vd)
  t2 = adjoint(U*u)
  rmul!(t1, Diagonal(1 ./ d))
  t1*t2
end

# Calculates (1 + UDVd)^-1, where U, D, Vd come from SVD decomp.
# More controlled handling of scales, however also slower.
# This is Yoni's `stable_invert_B`
# Seems to work even for large spread in D (i.e. low temperatures)!
function inv_one_plus_udv_scalettar(U,D,Vd)
  Dp = max.(D,1.)
  Dm = min.(D,1.)
  Dpinv = 1 ./ Dp

  l = copy(Vd')
  rmul!(l, Diagonal(Dpinv))

  r = copy(U)
  rmul!(r, Diagonal(Dm))

  u, d, vd = decompose_udv!(l+r)

  m = inv_udv(u,d,vd)
  lmul!(Diagonal(Dpinv), m)
  u, d, vd = decompose_udv!(m)

  mul!(m, Vd', u)
  # return m, d, vd
  rmul!(m, Diagonal(d))
  m*vd
end

# TODO: Optimize!
# I only made the function overwrite res. Otherwise it's unchanged compared to the one above.
function inv_one_plus_udv_scalettar!(mc, res, U,D,Vd)
  # all similars here could go into mc.s
  Dp = similar(D)
  Dm = similar(D)
  l = similar(Vd)
  l .= Vd'
  r = similar(U)
  r .= U
  # u, d, vd below could be preallocated, but does this matter for speed?
  # m could be preallocated, I guess

  Dp .= max.(D, 1)
  Dm .= min.(D, 1)

  Dp .\= 1 # Dp now Dpinv!!!

  rmul!(l, Diagonal(Dp))
  rmul!(r, Diagonal(Dm))

  u, d, vd = decompose_udv!(l + r)

  m = inv_udv(u,d,vd) # TODO: optimize
  lmul!(Diagonal(Dp), m)
  u, d, vd = decompose_udv!(m)

  mul!(m, Vd', u)
  # return m, d, vd
  rmul!(m, Diagonal(d))
  # res .= m*vd
  mul!(res, m, vd)
  nothing
end

function UDV_to_mat!(mat, U, D, Vd; invert=false)
    if !invert
        mat1 = copy(U)
        rmul!(mat1, Diagonal(D))
        mul!(mat, mat1, Vd)
    else #V D^(-1) Ud = (D^-1 *Vd)^(dagger) *Ud
        mat1 = copy(Vd)
        lmul!(Diagonal(1 ./ D), mat1)
        mul!(mat, mat1', U')
    end
    nothing
end

function UDV_to_mat(U, D, Vd; kw...)
  res = copy(U)
  UDV_to_mat!(res, U, D, Vd; kw...)
  res
end

# Calculates (UaDaVda + UbDbVdb)^-1
function inv_sum_udvs(Ua, Da, Vda, Ub, Db, Vdb)

    d=length(Da)

    # separating scales larger and smaller than unity
    Dap = max.(Da,1.)
    Dam = min.(Da,1.)
    Dbp = max.(Db,1.)
    Dbm = min.(Db,1.)

    # mat1 = Dam * Vda * Vdb' / Dbp

    mat1 = Vda * adjoint(Vdb)
    for j in 1:d, k in 1:d
        mat1[j,k]=mat1[j,k] * Dam[j]/Dbp[k]
    end

    # mat2 = 1/(Dap) * Ua' * Ub * Dbm
    mat2 = adjoint(Ua) * Ub
    for j in 1:d, k in 1:d
        mat2[j,k]=mat2[j,k] * Dbm[k]/Dap[j]
    end

    #mat1 = mat1 + mat2
    mat1 = mat1 + mat2

    # decompose mat1: U, D, Vd
    U, D, Vd = decompose_udv!(mat1)

    # invert inner part: mat1 = (U D Vd)^(-1) = mat1^(-1)
    UDV_to_mat!(mat1, U, D, Vd, invert=true)

    # mat1 = 1/Dbp * mat1 /Dap
    for j in 1:d, k in 1:d
        mat1[j,k]=mat1[j,k] / Dbp[j] / Dap[k]
    end

    #mat1 = U D Vd
    U, D, Vd = decompose_udv!(mat1)

    # U = Vdb' * U , Vd = Vd * Ua'
    mul!(mat1, Vdb', U)
    mul!(mat2, Vd, Ua')
    U = mat1
    Vd = mat2

    return U, D, Vd
end

# Calculates (UaDaVda + UbDbVdb)^-1
# TODO: Optimize!
# I only made the function overwrite res. Otherwise it's unchanged compared to the one above.
function inv_sum_udvs!(mc, res, Ua, Da, Vda, Ub, Db, Vdb)

    d=length(Da)

    #DXp = max (X%D, 1) ,DXm = min (X%D, 1) and similarly for Y.
    Dap = max.(Da,1.)
    Dam = min.(Da,1.)
    Dbp = max.(Db,1.)
    Dbm = min.(Db,1.)

    #mat1= DXm * X%Vd * (Y%Vd)^dagger /DYp
    mat1 = Vda * adjoint(Vdb)
    for j in 1:d, k in 1:d
        mat1[j,k]=mat1[j,k] * Dam[j]/Dbp[k]
    end

    #mat2 = 1/(DXp) * (X%U)^dagger * Y%U * DYm
    mat2 = adjoint(Ua) * Ub
    for j in 1:d, k in 1:d
        mat2[j,k]=mat2[j,k] * Dbm[k]/Dap[j]
    end

    #mat1 = mat1+mat2
    mat1 = mat1 + mat2

    #invert mat1: mat1=mat1^(-1)
    U, D, Vd = decompose_udv!(mat1)
    UDV_to_mat!(mat1, U, D, Vd, invert=true)

    #mat1 = 1/DYp * mat1 /DXp
    for j in 1:d, k in 1:d
        mat1[j,k]=mat1[j,k] / Dbp[j] / Dap[k]
    end

    #mat1 = U D Vd
    U, D, Vd = decompose_udv!(mat1)

    # U = (Y%Vd)^dagger * U , Vd = Vd * (X%U)^dagger
    mul!(mat1,adjoint(Vdb),U)
    mul!(mat2,Vd,adjoint(Ua))
    U=mat1
    Vd=mat2

    # return U, D, Vd
    UDV_to_mat!(res, U, D, Vd, invert=false)
    nothing
end


function bandec_trans(a::Array,n::Int64,m1::Int64,m2::Int64,al::Array,indx::Vector{Int64})
#      INTEGER m1,m2,mp,mpl,n,np,indx(n)
#      REAL d,a(np,mp),al(np,mpl),TINY
mm=m1+m2+1
@assert(size(a)  == (mm,n))
@assert(size(al) == (m1,n))
@assert(length(indx) == n)
TINY = one(eltype(a))*1.e-20
#      INTEGER i,j,k,l,mm
#      REAL dum
l=m1
czero = zero(eltype(a))
#dum = czero
for i=1:m1
  for j=m1+2-i:mm
    a[j-l,i]=a[j,i]
  end
  l -= 1
  for j=mm-l:mm
    a[j,i]=czero
  end
end
d = one(eltype(a))
l=m1
for k=1:n
  dum=a[1,k]
  i=k
  if (l < n)
    l += 1
  end
  for j=k+1:l
    if (abs(a[1,j]) > abs(dum))
      dum=a[1,j]
      i=j
    end
  end
  indx[k]=i
  if (dum == czero)
    a[1,k]=TINY
  end
  if (i != k)
    d = -d
    for j=1:mm
      dum=a[j,k]
      a[j,k]=a[j,i]
      a[j,i]=dum
    end
  end
  for i=k+1:l
    dum=a[1,i]/a[1,k]
    al[i-k,k]=dum
    for j=2:mm
      a[j-1,i] = a[j,i] - dum*a[j,k]
    end
    a[mm,i]=czero
  end
end
return d
end


using SumBassSolver
using Random

Random.seed!(121)
# Generate random data. For implementation input proper data
p = 0.003*rand(50)
q = 0.3*rand(50)
m = 1000 .+ 500*rand(50)
c0 = m .* rand(50) ./ 2
alpha = rand(50)
F = zeros(50)
T = fill(100.0,50)
Pi = 2500

x1,v1 = @time sumBassSolver(p,q,m,c0,alpha,F,T,Pi)
x2,v2 = @time sumBassSolver(p,q,m,c0,alpha,F,T,Pi,localopt=true)
x3,v3 = @time sumBassSolver(p,q,m,c0,alpha,F,T,Pi,alg="SB")
x4,v4 = @time sumBassSolver(p,q,m,c0,alpha,F,T,Pi,alg="SB",localopt=true)

vv1 = sum(alpha .* Bass.(x1,p,q,m,c0))
vv2 = sum(alpha .* Bass.(x2,p,q,m,c0))
vv3 = sum(alpha .* Bass.(x3,p,q,m,c0))
vv4 = sum(alpha .* Bass.(x4,p,q,m,c0))
sum(abs.(x1-x4))

lambdas = Bass_derivative.(x,p,q,m,c0).*alpha

i = findall(x -> abs(x-lambdas[1]) > 0.01,lambdas) # Sites at the boundary
x[i]