function decompose_hadamard()

    println(">>>>> Hadamard Gate <<<<<")

    return Dict{String, Any}(

    "num_qubits" => 2,
    "maximum_depth" => 3,
    "elementary_gates" => ["U3_1", "CNot_1_2", "Identity"],
    "target_gate" => QCOpt.get_full_sized_gate("H_1", 2),
    "objective" => "minimize_depth",
    "decomposition_type" => "exact",

    "U3_θ_discretization" => [0, π/2],
    "U3_ϕ_discretization" => [0, π/2],
    "U3_λ_discretization" => [0, π],
    )

end

function decompose_controlled_Z()

    println(">>>>> Controlled-Z Gate <<<<<")

    return Dict{String, Any}(

    "num_qubits" => 2,
    "maximum_depth" => 4,
    "elementary_gates" => ["U3_1", "U3_2", "CNot_1_2", "Identity"],
    "target_gate" => QCOpt.CZGate(),
    "objective" => "minimize_depth",
    "decomposition_type" => "exact",

    "U3_θ_discretization" => -π:π/2:π,
    "U3_ϕ_discretization" => -π:π/2:π,
    "U3_λ_discretization" => -π:π/2:π,
    )

end

function decompose_controlled_V()

    println(">>>>> Controlled-V Gate <<<<<")

    return Dict{String, Any}(

    "num_qubits" => 2,
    "maximum_depth" => 7,
    "elementary_gates" => ["H_1", "H_2", "T_1", "T_2", "Tdagger_1", "Tdagger_2", "CNot_1_2", "CNot_2_1", "Identity"],
    "target_gate" => QCOpt.CVGate(),
    "objective" => "minimize_depth",
    "decomposition_type" => "exact",
    )

end

function decompose_controlled_H()

    println(">>>>> Controlled-H Gate <<<<<")

    return Dict{String, Any}(

    "num_qubits" => 2,
    "maximum_depth" => 5,
    "elementary_gates" => ["U3_1", "U3_2", "CNot_1_2", "CNot_2_1", "Identity"],
    "target_gate" => QCOpt.CHGate(),
    "objective" => "minimize_depth",
    "decomposition_type" => "exact",

    "U3_θ_discretization" => -2*π:π/4:2*π,
    "U3_ϕ_discretization" => [0],
    "U3_λ_discretization" => [0],
    )

end

function decompose_controlled_H_with_R()

    println(">>>>> Controlled-H with R Gate <<<<<")

    return Dict{String, Any}(

    "num_qubits" => 2,
    "maximum_depth" => 5,
    "elementary_gates" => ["RY_1", "RY_2", "CNot_1_2", "Identity"],
    "target_gate" => QCOpt.CHGate(),
    "objective" => "minimize_depth",
    "decomposition_type" => "exact",

    "RY_discretization" => [-π/4, π/4, π/2, -π/2],
    "set_cnot_lower_bound" => 1,
    "set_cnot_upper_bound" => 2,
    )

end

function decompose_magic()

    println(">>>>> M Gate <<<<<")

    return Dict{String, Any}(

        "num_qubits" => 2,
        "maximum_depth" => 4,
        "elementary_gates" => ["U3_1", "U3_2", "CNot_2_1", "CNot_1_2", "Identity"],
        "target_gate" => QCOpt.MGate(),
        "objective" => "minimize_depth",
        "decomposition_type" => "exact",

        "U3_θ_discretization" => -π:π/2:π,
        "U3_ϕ_discretization" => -π:π/2:π,
        "U3_λ_discretization" => -π:π/2:π,
        )

end

function decompose_magic_using_CNOT_1_2()

    println(">>>>> M Gate <<<<<")

    return Dict{String, Any}(

        "num_qubits" => 2,
        "maximum_depth" => 5,
        "elementary_gates" => ["U3_1", "U3_2", "CNot_1_2", "Identity"],
        "target_gate" => QCOpt.MGate(),
        "objective" => "minimize_depth",
        "decomposition_type" => "exact",

        "U3_θ_discretization" => -π:π/2:π,
        "U3_ϕ_discretization" => -π:π:π,
        "U3_λ_discretization" => -π:π/2:π,
        )

end

function decompose_magic_using_SHCnot()

    println(">>>>> M gate using S, H and CNOT Gate <<<<<")

    return Dict{String, Any}(

        "num_qubits" => 2,
        "maximum_depth" => 5,
        "elementary_gates" => ["S_1", "S_2", "H_1", "H_2", "CNot_1_2", "CNot_2_1", "Identity"],
        "target_gate" => QCOpt.MGate(),
        "objective" => "minimize_depth",
        "decomposition_type" => "exact",
        )

end

function decompose_S()

    println(">>>>> S Gate <<<<<")

    return Dict{String, Any}(

        "num_qubits" => 2,
        "maximum_depth" => 3,
        "elementary_gates" => ["U3_1", "U3_2", "CNot_1_2", "Identity"],
        "target_gate" => QCOpt.get_full_sized_gate("S_1", 2),
        "objective" => "minimize_depth",
        "decomposition_type" => "exact",

        "U3_θ_discretization" => [0, π/4, π/2],
        "U3_ϕ_discretization" => [-π/2, 0, π/2],
        "U3_λ_discretization" => [0, π],
        )
end

function decompose_revcnot()

    println(">>>>> Reverse CNOT Gate <<<<<")

    return Dict{String, Any}(

    "num_qubits" => 2,
    "maximum_depth" => 5,
    "elementary_gates" => ["H_1", "H_2", "CNot_1_2", "Identity"],
    "target_gate" => QCOpt.CNotRevGate(),
    "objective" => "minimize_depth",
    "decomposition_type" => "exact",
    )

end

function decompose_revcnot_with_U()

    println(">>>>> Reverse CNOT using U3 and CNOT Gates <<<<<")

    return Dict{String, Any}(

        "num_qubits" => 2,
        "maximum_depth" => 5,
        "elementary_gates" => ["U3_1", "U3_2", "CNot_1_2", "Identity"],
        "target_gate" => QCOpt.CNotRevGate(),
        "objective" => "minimize_depth",
        "decomposition_type" => "exact",

        "U3_θ_discretization" => -π:π/4:π,
        "U3_ϕ_discretization" => [0],
        "U3_λ_discretization" => [0],
        )

end

function decompose_swap()

    println(">>>>> SWAP Gate <<<<<")

    return Dict{String, Any}(

        "num_qubits" => 2,
        "maximum_depth" => 5,
        "elementary_gates" => ["CNot_2_1", "CNot_1_2", "Identity"],
        "target_gate" => QCOpt.SwapGate(),
        "objective" => "minimize_depth",
        "decomposition_type" => "exact",
        )

end

function decompose_W()

    println(">>>>> W Gate <<<<<")

    return Dict{String, Any}(

        "num_qubits" => 2,
        "maximum_depth" => 5,
        "elementary_gates" => ["U3_2", "CNot_2_1", "CNot_1_2", "Identity"],
        "target_gate" => QCOpt.WGate(),
        "objective" => "minimize_depth",
        "decomposition_type" => "exact",

        "U3_θ_discretization" => [-π/4, 0, π/4],
        "U3_ϕ_discretization" => [0, π/4],
        "U3_λ_discretization" => [0, π/2],
        )

end

function decompose_W_using_HCnot()

    println(">>>>> W Gate using H and CNOT gates <<<<<")

    return Dict{String, Any}(

        "num_qubits" => 2,
        "maximum_depth" => 6,
        "elementary_gates" => ["CH_1_2", "CNot_2_1", "CNot_1_2", "Identity"],
        "target_gate" => QCOpt.WGate(),
        "objective" => "minimize_depth",
        "decomposition_type" => "exact",
        )

end

function decompose_HCoinGate()

    println(">>>>> Hadamard Coin gate <<<<<")

    return Dict{String, Any}(

        "num_qubits" => 2,
        "maximum_depth" => 10,
        "elementary_gates" => ["Y_1", "Y_2", "Z_1", "Z_2", "T_2", "Tdagger_1", "Sdagger_1", "SX_1", "SXdagger_2", "CNot_2_1", "CNot_1_2", "Identity"],
        "target_gate" => QCOpt.HCoinGate(),
        "objective" => "minimize_depth",
        "decomposition_type" => "exact",
        )

end

function decompose_GroverDiffusion_using_HX()

    println(">>>>> Grover's Diffusion Operator <<<<<")

    return Dict{String, Any}(

        #Reference: https://arxiv.org/pdf/1804.03719.pdf (Fig 6)

        "num_qubits" => 2,
        "maximum_depth" => 10,
        "elementary_gates" => ["X_1", "X_1xX_2", "H_1xH_2", "X_2", "H_1", "H_2", "CNot_1_2", "Identity"],
        "target_gate" => QCOpt.GroverDiffusionGate(),
        "objective" => "minimize_depth",
        "decomposition_type" => "exact",
        )

end

function decompose_GroverDiffusion_using_Clifford()

    println(">>>>> Grover's Diffusion Operator <<<<<")

    return Dict{String, Any}(

        "num_qubits" => 2,
        "maximum_depth" => 6,
        "elementary_gates" => ["X_1", "X_2", "H_1", "H_2", "S_1", "S_2", "T_1", "T_2", "Y_1", "Y_2", "CNot_1_2", "Identity"],
        "target_gate" => QCOpt.GroverDiffusionGate(),
        "objective" => "minimize_depth",
        "decomposition_type" => "exact",
        )

end

function decompose_GroverDiffusion_using_U3()

    println(">>>>> Grover's Diffusion Operator using U3 gate <<<<<")

    return Dict{String, Any}(

        "num_qubits" => 2,
        "maximum_depth" => 10,
        "elementary_gates" => ["U3_1", "CNot_1_2", "Identity"],
        "target_gate" => QCOpt.GroverDiffusionGate(),
        "objective" => "minimize_depth",
        "decomposition_type" => "exact",

        "U3_θ_discretization" => -π:π/2:π,
        "U3_ϕ_discretization" => -π:π/2:π,
        "U3_λ_discretization" => [0],
        )

end

function decompose_iSwap()

    println(">>>>> iSwap Gate <<<<<")

    return Dict{String, Any}(

        "num_qubits" => 2,
        "maximum_depth" => 10,
        "elementary_gates" => ["T_1", "T_2", "Tdagger_1", "Tdagger_2", "H_1", "H_2", "CNot_1_2", "CNot_2_1", "Identity"],
        # "elementary_gates" => ["S_1", "S_2", "Sdagger_1", "Sdagger_2", "H_1", "H_2", "CNot_1_2", "CNot_2_1", "Identity"],
        "target_gate" => QCOpt.iSwapGate(),
        "objective" => "minimize_depth",
        "decomposition_type" => "exact",
        )

end

function decompose_qft2_using_R()
    println(">>>>> QFT2 Gate using R and CNOT gates <<<<<")

    return Dict{String, Any}(

        "num_qubits" => 2,
        "maximum_depth" => 10,
        "elementary_gates" => ["RX_1", "RZ_1", "RZ_2", "CNot_1_2", "CNot_2_1", "Identity"],
        "target_gate" => QCOpt.QFT2Gate(),
        "objective" => "minimize_depth",
        "decomposition_type" => "approximate",

        "RX_discretization" => [π/2],
        "RZ_discretization" => [-π/4, π/2, 3*π/4, 7*π/4],
        )

end

function decompose_qft2_using_HT()
    println(">>>>> QFT2 Gate using H, T, CNOT gates <<<<<")

    return Dict{String, Any}(

        "num_qubits" => 2,
        "maximum_depth" => 8,
        "elementary_gates" => ["H_1", "H_2", "T_1", "T_2", "Tdagger_1", "Tdagger_2", "CNot_1_2", "CNot_2_1", "Identity"],
        "target_gate" => QCOpt.QFT2Gate(),
        "objective" => "minimize_depth",
        "decomposition_type" => "exact",
        )

end


struct SO3 <: Manifold end

dim(::Type{SO3}) = 3
embdim(::Type{SO3}) = 9

abstract type ExpMap <: ChartID end
isglobal(::Chart{<:ExpMap,SO3}) = false
struct ExpMap0 <: ExpMap end
struct ExpMapX <: ExpMap end
struct ExpMapY <: ExpMap end
struct ExpMapZ <: ExpMap end

skew(w) = SA[0.   -w[3]  w[2]
             w[3]  0.   -w[1]
            -w[2]  w[1]  0.]
unskew(ŵ) = SA[ŵ[3,2], ŵ[1,3], ŵ[2,1]]

function logSO3(R)
    trR = tr(R)
    if trR < 3
        θ = acos((trR-1)/2)
        # ŵ = (R - R')/(2*sin(θ))
        # w = unskew(ŵ)
        w = SA[R[3,2] - R[2,3], R[1,3] - R[3,1], R[2,1] - R[1,2]]/(2*sin(θ))
        return θ*w
    else
        return @SVector zeros(3)
    end
end

function expSO3(w)
    w == (@SVector zeros(3)) && (@show "expSO3 singularity"; return SMatrix{3,3,eltype(w)}(I))
    ŵ = skew(w)
    wnorm = norm(w)
    SMatrix{3,3,eltype(w)}(I) + ŵ*sin(wnorm)/wnorm + ŵ^2*(1-cos(wnorm))/wnorm^2
end

tr_vec(pe::SVector{9,S}) where S = pe[1] + pe[5] + pe[9]
chart_boundary_test(pe, ::Chart{ExpMap0,SO3}) = tr_vec(pe)
chart_boundary_test(pe, ::Chart{ExpMapX,SO3}) = tr(RotX(π/2)*reshape(pe, Size(3,3)))
chart_boundary_test(pe, ::Chart{ExpMapY,SO3}) = tr(RotY(π/2)*reshape(pe, Size(3,3)))
chart_boundary_test(pe, ::Chart{ExpMapZ,SO3}) = tr(RotZ(π/2)*reshape(pe, Size(3,3)))

chart_to_emb(p0, ::Chart{ExpMap0,SO3}) = reshape(expSO3(p0), Size(9))
chart_to_emb(px, ::Chart{ExpMapX,SO3}) = reshape(RotX(-π/2)*expSO3(px), Size(9))
chart_to_emb(py, ::Chart{ExpMapY,SO3}) = reshape(RotY(-π/2)*expSO3(py), Size(9))
chart_to_emb(pz, ::Chart{ExpMapZ,SO3}) = reshape(RotZ(-π/2)*expSO3(pz), Size(9))

emb_to_chart(pe, ::Chart{ExpMap0,SO3}) = logSO3(reshape(pe, Size(3,3)))
emb_to_chart(pe, ::Chart{ExpMapX,SO3}) = logSO3(RotX(π/2)*reshape(pe, Size(3,3)))
emb_to_chart(pe, ::Chart{ExpMapY,SO3}) = logSO3(RotY(π/2)*reshape(pe, Size(3,3)))
emb_to_chart(pe, ::Chart{ExpMapZ,SO3}) = logSO3(RotZ(π/2)*reshape(pe, Size(3,3)))

function zero_jacobian_component(::Chart{ExpMap0,SO3})
    SA[0 0  0  0 0 1 0 -1 0
       0 0 -1  0 0 0 1  0 0
       0 1  0 -1 0 0 0  0 0.]
end

function chart_to_emb_jacobian(p, C::Chart{ExpMap0,SO3})
    if norm(p - (@SVector zeros(3))) < 1e-12
        return zero_jacobian_component(C)'
    else
        return ForwardDiff.jacobian(p -> chart_to_emb(p, C), p)
    end
end

function emb_to_chart_jacobian(pe, C::Chart{ExpMap0,SO3})
    if (chart_boundary_test(pe, C) - 3) < 1e-12
        return zero_jacobian_component(C)/2
    else
        return ForwardDiff.jacobian(pe -> emb_to_chart(pe, C), pe)
    end
end

chart_choice_coord_rep(::Chart{<:ExpMap,SO3}) = EmbRep()
function choose_chart_emb(::EmbRep, pe, Cn::Chart{<:ExpMap,SO3})
    C0, Cx = Chart{ExpMap0,SO3}(), Chart{ExpMapX,SO3}()
    Cy, Cz = Chart{ExpMapY,SO3}(), Chart{ExpMapZ,SO3}()

    lo_bound, hi_bound = -0.5, 2.5
    val = chart_boundary_test(pe, Cn)
    if val < lo_bound || val > hi_bound
        val0 = chart_boundary_test(pe, C0)
        valx = chart_boundary_test(pe, Cx)
        valy = chart_boundary_test(pe, Cy)
        valz = chart_boundary_test(pe, Cz)

        val0 = val0 > hi_bound ? -1. : val0
        valx = valx > hi_bound ? -1. : valx
        valy = valy > hi_bound ? -1. : valy
        valz = valz > hi_bound ? -1. : valz

        val_max, ind_max = findmax(SA[val0, valx, valy, valz])
        return (C0, Cx, Cy, Cz)[ind_max]
    else
        return Cn
    end
end

function apply_rotation_emb(xme, R)
    Rtarget = reshape(xme, Size(3,3))
    reshape(R*Rtarget, Size(9))
end

function apply_inverse_rotation_jacobian_emb(xme)
    R = reshape(xme, Size(3,3))
    Rinv = R'
    ForwardDiff.jacobian(xme -> apply_rotation_emb(xme, Rinv), xme)
end

function apply_rotation_chart(xm, R, Ci::Chart{<:ExpMap,SO3})
    xme = chart_to_emb(xm, Ci)
    xme_out = apply_rotation_emb(xme, R)
    xm_out = emb_to_chart(xme_out, Ci)
end

function apply_rotation_chart_emb(xm, R, Ci::Chart{<:ExpMap,SO3})
    xme = chart_to_emb(xm, Ci)
    xme_out = apply_rotation_emb(xme, R)
end

function apply_inverse_rotation_jacobian_chart(xm, Ci::Chart{<:ExpMap,SO3})
    xme = chart_to_emb(xm, Ci)
    R = reshape(xme, Size(3,3))
    Rinv = R'
    xme_out = reshape(SMatrix{3,3,Float64}(I), Size(9))
    J1 = ForwardDiff.jacobian(xm -> apply_rotation_chart_emb(xm, Rinv, Ci), xm)
    J2 = emb_to_chart_jacobian(xme_out, Ci)
    J2*J1
end

function default_metric(pn, task::Task{<:TaskMap{M,SO3}}, CN::Chart{<:ExpMap,SO3}) where M
    C0 = Chart{ExpMap0,SO3}()
    JL = apply_inverse_rotation_jacobian_chart(pn, CN)
    Jϕ = chart_transition_jacobian(pn, CN, C0)
    J = Jϕ*JL
    J'*SMatrix{3,3,Float64}(I)*J
end

default_weight_metric(pn, wn, task::Task{<:TaskMap{M,SO3}}, CN::Chart{<:ExpMap,SO3}) where M =
    default_metric(pn, task, CN)


using COSMO
using Random, Test, LinearAlgebra

verbosity = 0 # Change to 2 to see LOBPCG's output
rng = MersenneTwister(1)
λ = [1:6; -50.00:-0.00]
n = length(λ)
Q = Matrix(qr(randn(rng, n, n)).Q)
X = Q*Diagonal(λ)*Q'

n_ = Int(n*(n + 1)/2)

function exact_projection(A::Matrix)
    l, V = eigen(Symmetric(A))
    return Symmetric(V*Diagonal(max.(l, 0))*V')
end

function lobpcg_unit_test!(cone, X::Matrix)
    x = Main.COSMO.extract_upper_triangle(X)
    X_proj_exact = exact_projection(X)
    COSMO.project!(x, cone)
    X_proj = Main.COSMO.populate_upper_triangle(x)

    @test norm(X_proj - X_proj_exact) < 1e-5 # This is tolerance is hand-wavy
end

function lobpcg_unit_tests(X::Matrix)
    cone = Main.COSMO.PsdConeTriangleLOBPCG(n_, verbosity=verbosity, max_iter=200)
    @testset "Initial call (exact eigendecomposition)" begin
        lobpcg_unit_test!(cone, X)
        @test cone.exact_projections == 1
        @test cone.lobpcg_iterations == 0
    end

    @testset "Warm starting - already solved to tolerance" begin
        cone.iteration = 10^3
        lobpcg_unit_test!(cone, X + 1e-6*Symmetric(randn(rng, n, n)))
        @test cone.exact_projections == 1
        @test cone.lobpcg_iterations == 1
    end
    @testset "Warm starting - a few iterations required" begin
        cone.iteration = 10^8
        lobpcg_unit_test!(cone, X + 1e-6*Symmetric(randn(rng, n, n)))
        @test cone.exact_projections == 1
        @test cone.lobpcg_iterations >= 3
        @test cone.lobpcg_iterations < 10
    end

    @testset "Warm starting - expand subspace" begin
        cone.iteration = 10^8
        if cone.is_subspace_positive
            lobpcg_unit_test!(cone, X + 3.1*I)
        else
            lobpcg_unit_test!(cone, X - 3.1*I)
        end
        @test size(cone.U, 2) > 6
        @test cone.exact_projections == 1
        @test cone.lobpcg_iterations >= 3
        @test cone.lobpcg_iterations < 50
    end

    @testset "Warm starting - contract subspace" begin
        cone.iteration = 10^8
        if cone.is_subspace_positive
            lobpcg_unit_test!(cone, X - 1.1*I)
        else
            lobpcg_unit_test!(cone, X + 1.1*I)
        end
        @test size(cone.U, 2) < 6
        @test cone.exact_projections == 1
        @test cone.lobpcg_iterations >= 3
        @test cone.lobpcg_iterations < 60
    end
end

@testset "LOBPCG PSD Projection" begin
    @testset "Track Positive Eigenspace" begin
        lobpcg_unit_tests(X)
    end
    @testset "Track Negative Eigenspace" begin
        lobpcg_unit_tests(-X)
    end
end

nothing


# Uses method from Parviainen et al. 2020 to compute the sky-separation.

"""Compute the derivative components of the series expansion."""
function components(d::AbstractVector{T},h::T) where T<:AbstractFloat
    A = d[7]-d[1]; B = d[1]+d[7]; C = d[2]-d[6]
    D = d[2]+d[6]; E = d[5]-d[3]; F = d[3]+d[5]
    v = dot(SVector(A, C, E), COEFF.vd)/h
    a = dot(SVector(B, D, F, d[4]), COEFF.ad)/(h*h)
    j = dot(SVector(-A, -C, -E), COEFF.jd)/(h*h*h)
    s = dot(SVector(B, D, F , d[4]), COEFF.sd)/(h*h*h*h)
    return SVector(d[4],v,0.5*a,(1/6)*j,(1/24)*s)
end

"""Compute the impact parameter at t=tc using a series expansion about t0."""
function compute_impact_parameter(tc::T, t0::T, xc::AbstractVector{T}, yc::AbstractVector{T}) where T<:Real
    t = tc - t0
    ts = SVector(1.0, t, t*t, t*t*t, t*t*t*t)
    lx = dot(xc, ts)
    ly = dot(yc, ts)
    return sqrt(lx*lx + ly*ly)
end

function compute_impact_parameter(tc::T,t0::T,h::T,points::AbstractMatrix{T}) where T<:Real
    t = tc - t0
    ts = SVector(1.0,t,t*t,t*t*t,t*t*t*t)
    xc = components(points[:,1], h)
    yc = components(points[:,2], h)
    lx = dot(xc,ts)
    ly = dot(yc,ts)
    return sqrt(lx*lx + ly*ly)
end

function compute_impact_parameter!(tc::T, t0::T, xc, yc, dxc, dyc, grad) where T<:Real
    t = tc - t0
    ts = SVector(1.0, t, t*t, t*t*t, t*t*t*t)
    lx = dot(xc, ts)
    ly = dot(yc, ts)
    b = sqrt(lx*lx + ly*ly)

    dbdlx = lx/b; dbdly = ly/b;
    N = length(grad)
    for p in 1:N
        dlxdp = dot(dxc[p],ts)
        dlydp = dot(dyc[p],ts)
        grad[p] = dbdlx*dlxdp + dbdly*dlydp
    end

    return b
end


include("CenterOfMass.jl")

np = 5
thetas = 2*pi * (1/np:1/np:1)
circle_ps = map(theta->[cos(theta);sin(theta);0], thetas)

na = 5
thetas = 2*pi * (1/na:1/na:1)
r_attach = 0.8
attachment_ps = map(theta->r_attach*[cos(theta);sin(theta);0], thetas)

get_com() = r_attach * rand_in_circle()

com = get_com()

lin_prog, lin_fs = critical_force_from_points(circle_ps, com, attachment_ps[1], 1)
iterative, it_fs = critical_force_from_points(circle_ps, com, attachment_ps[1], 1, critical_force_iterative)

println("Lin prog solution: $(lin_prog), iterative solution $(iterative)")


"""
    rCIP.jl

This is a Julia port of a MATLAB implementation of batch and incremental
(Renyi's) representative Cross Information Potential (rCIP) Cluster Validity Index.

Authors:
MATLAB implementation: Leonardo Enzo Brito da Silva
Julia port: Sasha Petrenko <sap625@mst.edu>

REFERENCES
[1] L. E. Brito da Silva, N. M. Melton, and D. C. Wunsch II, "Incremental
Cluster Validity Indices for Hard Partitions: Extensions  and  Comparative
Study," ArXiv  e-prints, Feb 2019, arXiv:1902.06711v1 [cs.LG].
[2] E. Gokcay and J. C. Principe, "A new clustering evaluation function
using Renyi’s information potential," in Proc. Int. Conf. Acoust., Speech,
Signal Process. (ICASSP), vol. 6. Jun. 2000, pp. 3490–3493.
[3] E. Gokcay and J. C. Principe, "Information theoretic clustering," IEEE
Trans. Pattern Anal. Mach. Intell., vol. 24, no. 2, pp. 158–171, Feb. 2002.
[4] D. Araújo, A. D. Neto, and A. Martins, "Representative cross information
potential clustering," Pattern Recognit. Lett., vol. 34, no. 16,
pp. 2181–2191, Dec. 2013.
[5] D. Araújo, A. D. Neto, and A. Martins, "Information-theoretic clustering:
A representative and evolutionary approach," Expert Syst. Appl.,
vol. 40, no. 10, pp. 4190–4205, Aug. 2013.
[6] R. O. Duda, P. E. Hart, and D. G. Stork, Pattern Classification, 2nd ed.
John Wiley & Sons, 2000.
"""

"""
    rCIP

The stateful information of the (Renyi's) representative Cross Information Potential (rCIP) Cluster Validity Index.

# References
1. L. E. Brito da Silva, N. M. Melton, and D. C. Wunsch II, "Incremental Cluster Validity Indices for Hard Partitions: Extensions  and  Comparative Study," ArXiv  e-prints, Feb 2019, arXiv:1902.06711v1 [cs.LG].
2. E. Gokcay and J. C. Principe, "A new clustering evaluation function using Renyi’s information potential," in Proc. Int. Conf. Acoust., Speech, Signal Process. (ICASSP), vol. 6. Jun. 2000, pp. 3490–3493.
3. E. Gokcay and J. C. Principe, "Information theoretic clustering," IEEE Trans. Pattern Anal. Mach. Intell., vol. 24, no. 2, pp. 158–171, Feb. 2002.
4. D. Araújo, A. D. Neto, and A. Martins, "Representative cross information potential clustering," Pattern Recognit. Lett., vol. 34, no. 16, pp. 2181–2191, Dec. 2013.
5. D. Araújo, A. D. Neto, and A. Martins, "Information-theoretic clustering: A representative and evolutionary approach," Expert Syst. Appl., vol. 40, no. 10, pp. 4190–4205, Aug. 2013.
6. R. O. Duda, P. E. Hart, and D. G. Stork, Pattern Classification, 2nd ed. John Wiley & Sons, 2000.
"""
mutable struct rCIP <: CVI
    label_map::LabelMap
    dim::Integer
    n_samples::Integer
    n::IntegerVector            # dim
    v::RealMatrix               # dim x n_clusters
    sigma::Array{Float, 3}     # dim x dim x n_clusters
    constant::Float
    D::RealMatrix               # n_clusters x n_clusters
    delta_term::RealMatrix      # dim x dim
    n_clusters::Integer
    criterion_value::Float
end # rCIP <: CVI

"""
    rCIP()

Default constructor for the (Renyi's) representative Cross Information Potential (rCIP) Cluster Validity Index.
"""
function rCIP()
    rCIP(
        LabelMap(),                         # label_map
        0,                                  # dim
        0,                                  # n_samples
        Array{Integer, 1}(undef, 0),        # n
        Array{Float, 2}(undef, 0, 0),       # v
        Array{Float, 3}(undef, 0, 0, 0),    # sigma
        0.0,                                # constant
        Array{Float, 2}(undef, 0, 0),       # D
        Array{Float, 2}(undef, 0, 0),       # delta_term
        0,                                  # n_clusters
        0.0                                 # criterion_value
    )
end # rCIP()

"""
    setup!(cvi::rCIP, sample::Vector{T}) where {T<:RealFP}
"""
function setup!(cvi::rCIP, sample::Vector{T}) where {T<:RealFP}
    # Get the feature dimension
    cvi.dim = length(sample)
    # Initialize the 2-D and 3-D arrays with the correct feature dimension
    cvi.v = Array{T, 2}(undef, cvi.dim, 0)
    cvi.sigma = Array{T, 3}(undef, cvi.dim, cvi.dim, 0)
    # Calculate the delta term
    epsilon = 12
    delta = 10^(-epsilon/cvi.dim)
    cvi.delta_term = Matrix{Float64}(LinearAlgebra.I, cvi.dim, cvi.dim).*delta
end # setup!(cvi::rCIP, sample::Vector{T}) where {T<:RealFP}

function param_inc!(cvi::rCIP, sample::RealVector, label::Integer)
    # Get the internal label
    i_label = get_internal_label!(cvi.label_map, label)

    n_samples_new = cvi.n_samples + 1
    if isempty(cvi.v)
        setup!(cvi, sample)
    end
    cvi.constant = 1/sqrt((2*pi)^cvi.dim)

    if i_label > cvi.n_clusters
        n_new = 1
        v_new = sample
        sigma_new = cvi.delta_term
        if cvi.n_clusters == 0
            D_new = zeros(1, 1)
        else
            D_new = zeros(cvi.n_clusters + 1, cvi.n_clusters + 1)
            D_new[1:cvi.n_clusters, 1:cvi.n_clusters] = cvi.D
            d_column_new = zeros(cvi.n_clusters + 1)
            for jx = 1:cvi.n_clusters
                diff_m = v_new - cvi.v[:, jx]
                sigma_q = sigma_new + cvi.sigma[:,:,jx]
                d_column_new[jx] = cvi.constant * (1/sqrt(det(sigma_q)))*exp(-0.5*transpose(diff_m)*inv(sigma_q)*diff_m)
                # d_column_new[jx] = sum((v_new - cvi.v[:, jx]).^2)
            end
            D_new[:, i_label] = d_column_new
            D_new[i_label, :] = transpose(d_column_new)
        end
        # Update 1-D parameters with a push
        cvi.n_clusters += 1
        push!(cvi.n, n_new)
        # Update 2-D parameters with appending and reassignment
        cvi.v = [cvi.v v_new]
        cvi.D = D_new
        cvi.sigma = cat(cvi.sigma, sigma_new, dims=3)
    else
        n_new = cvi.n[i_label] + 1
        v_new = (1 - 1/n_new) .* cvi.v[:, i_label] + (1/n_new) .* sample
        diff_x_v = sample - cvi.v[:, i_label]
        if n_new > 1
            sigma_new = ((n_new - 2)/(n_new - 1))*(cvi.sigma[:,:,i_label] - cvi.delta_term) +
                (1/n_new)*(diff_x_v*transpose(diff_x_v)) + cvi.delta_term
        else
            sigma_new = cvi.delta_term
        end
        d_column_new = zeros(cvi.n_clusters)
        for jx = 1:cvi.n_clusters
            if jx == i_label
                continue
            end
            diff_m = v_new - cvi.v[:, jx]
            sigma_q = sigma_new + cvi.sigma[:,:,jx]
            d_column_new[jx] = cvi.constant*(1/sqrt(det(sigma_q)))*exp(-0.5*transpose(diff_m)*inv(sigma_q)*diff_m)
        end
        # Update parameters
        cvi.n[i_label] = n_new
        cvi.v[:, i_label] = v_new
        cvi.sigma[:,:,i_label] = sigma_new
        cvi.D[:, i_label] = d_column_new
        cvi.D[i_label, :] = transpose(d_column_new)
    end
    cvi.n_samples = n_samples_new
end # param_inc!(cvi::rCIP, sample::RealVector, label::Integer)

function param_batch!(cvi::rCIP, data::RealMatrix, labels::IntegerVector)
    cvi.dim, cvi.n_samples = size(data)
    cvi.constant = 1/sqrt((2*pi)^cvi.dim)
    # Calculate the delta term
    epsilon = 12
    delta = 10^(-epsilon/cvi.dim)
    cvi.delta_term = Matrix{Float64}(LinearAlgebra.I, cvi.dim, cvi.dim).*delta

    # Take the average across all samples, but cast to 1-D vector
    u = unique(labels)
    cvi.n_clusters = length(u)
    cvi.n = zeros(Integer, cvi.n_clusters)
    cvi.v = zeros(cvi.dim, cvi.n_clusters)
    cvi.sigma = zeros(cvi.dim, cvi.dim, cvi.n_clusters)
    cvi.D = zeros(cvi.n_clusters, cvi.n_clusters)

    for ix = 1:cvi.n_clusters
        subset = data[:, findall(x->x==u[ix], labels)]
        cvi.n[ix] = size(subset, 2)
        cvi.v[1:cvi.dim, ix] = mean(subset, dims=2)
        if cvi.n[ix] > 1
            cvi.sigma[:,:,ix] = (1/(cvi.n[ix] - 1)) *
                ((subset*transpose(subset)) - cvi.n[ix].*cvi.v[:,ix]*transpose(cvi.v[:,ix])) +
                cvi.delta_term
        else
            cvi.sigma[:,:,ix] = cvi.delta_term
        end
    end
    for ix = 1 : (cvi.n_clusters - 1)
        for jx = ix + 1 : cvi.n_clusters
            diff_m = cvi.v[:, ix] - cvi.v[:, jx]
            sigma_q = cvi.sigma[:,:,ix] + cvi.sigma[:,:,jx]
            cvi.D[jx, ix] = cvi.constant*(1/sqrt(det(sigma_q)))*exp(-0.5*transpose(diff_m)*inv(sigma_q)*diff_m)
        end
    end
    cvi.D = cvi.D + transpose(cvi.D)
end # param_batch!(cvi::rCIP, data::RealMatrix, labels::IntegerVector)

function evaluate!(cvi::rCIP)
    # Assume a symmetric dimension
    dim = size(cvi.D)[1]
    if dim > 1
        values = [cvi.D[i,j] for i = 1:dim, j=1:dim if j > i]
        cvi.criterion_value = sum(values)
    else
        cvi.criterion_value = 0.0
    end
end # evaluate(cvi::rCIP)


using TomoForward
using Images
using Plots
using XfromProjections
using StaticArrays
#using PyCall
using Logging
using Printf
using Suppressor

#include("../../phantoms/simple_phantoms.jl")
#include("../tv_primaldual_flow.jl")
#include("../optical_flow.jl")
include("../../phantoms/sperm_phantom.jl")

H,W = 609, 609
detmin, detmax = -38.0, 38.0
function radon_operator(height, width, detcount, angles)
    proj_geom = ProjGeom(0.5, detcount, angles)
    A = fp_op_parallel2d_line(proj_geom, height, width, -38.0,38.0, -38.0,38.0)
    return A
end

#Create phantom at smaller gridsize than reconstruction
grid = collect(detmin:0.1:detmax)
r(s) = 1.0
frames, tracks = get_sperm_phantom(2,r,grid)
detcount = Int(floor(H*1.4))
#match size of input image (generating data)
nangles = 4
a = range(0.0,2*pi+pi/4, length=nangles*size(frames)[3])#rand(0.0:0.001:π, nangles*size(frames)[3])
angles = map(t -> a[(t-1)*nangles+1:(t-1)*nangles+nangles], 1:size(frames)[3])
As = map(t -> radon_operator(size(frames[:,:,1],1),size(frames[:,:,1],2),detcount, angles[t]),1:size(frames)[3])

bsex = zeros(size(As[1])[1],size(frames)[3])

map(t -> bsex[:,t] = As[t]*vec(frames[:,:,t]), 1:size(frames)[3])
@info "adding gaussion noise at level 0.01"
rho = 0.01
e = randn(size(bsex));
e = rho*norm(bsex)*e/norm(e);
bs = bsex + e;

u0s = zeros(H,W,size(frames)[3])

As  = map(t -> radon_operator(H,W, detcount, angles[t]),1:size(u0s)[3])
cwd = @__DIR__
path = normpath(joinpath(@__DIR__, "result"))
!isdir(path) && mkdir(path)
cd(path)

u0s = zeros(H,W,size(frames)[3])
mask = ones(size(u0s)...)

c=1.1
w_tvs = [0.0001,0.001,0.01,0.05,0.1,0.2,0.3,0.5]
w_flows = [0.0001,0.001,0.01,0.05,0.1,0.2,0.3,0.5]
prog = Progress(size(w_tvs,1)*size(w_flows,1),60, "Testing parameters")

for w_tv = w_tvs
    niter=450
    us_tv = zeros(H,W,size(frames)[3])
    for t = 1:size(frames)[3]
        A = As[t]
        p = bs[:,t]
        u0 = u0s[:,:,t]
        @suppress begin
            us_tv[:,:,t] .= recon2d_tv_primaldual!(us_tv[:,:,t], A, p, niter, w_tv, c)
        end
    end
    for w_flow = w_flows
        niter=50
        @suppress begin
            us_flow = recon2d_tv_primaldual_flow(As, bs, deepcopy(us_tv), niter, 200, w_tv, w_flow, c, mask, 0.012)
            filenamepng = @sprintf "flow_2ms_%f_%f__3.png" w_tv w_flow;
            #PLOT
            plot(Gray.(us_flow[:,:,1]), aspect_ratio=:equal, framestyle=:none, legend=false, yflip=false)
            plot!(tracks[1][:,1], tracks[1][:,2], aspect_ratio=:equal, linewidth=5)
            savefig(filenamepng)
            next!(prog)
        end
    end
end

# angles = rand(1:0.001:π,3)
# #Create phantom at smaller gridsize than reconstruction
# frames, tracks = get_sperm_phantom(2,r, grid)
# detcount = Int(floor(H*1.4))
# #match size of input image (generating data)
# As = map(t -> radon_operator(size(frames[:,:,1],1),size(frames[:,:,1],2),detcount, angles[t]),1:size(frames)[3])
# bsex = zeros(size(As[1])[1],size(frames)[3])
#
# map(t -> bsex[:,t] = As[t]*vec(frames[:,:,t]), 1:size(frames)[3])
# @info "adding gaussion noise at level 0.01"
# rho = 0.01
# e = randn(size(bsex));
# e = rho*norm(bsex)*e/norm(e);
# bs = bsex + e;
#
# As  = map(t -> radon_operator(H,W, detcount, angles[t]),1:size(u0s)[3])
#
# for w_tv = w_tvs
#     niter=450
#     us_tv = zeros(H,W,size(frames)[3])
#     for t = 1:size(frames)[3]
#         A = As[t]
#         p = bs[:,t]
#         u0 = u0s[:,:,t]
#         @suppress begin
#             us_tv[:,:,t] .= recon2d_tv_primaldual!(us_tv[:,:,t], A, p, niter, w_tv, c)
#         end
#     end
#     for w_flow = w_flows
#         niter=50
#         @suppress begin
#             us_flow = recon2d_tv_primaldual_flow(As, bs, deepcopy(us_tv), niter, 20, w_tv, w_flow, c)
#             filenamepng = @sprintf "flow_radn_ang_2ms_%f_%f_.png" w_tv w_flow;
#             #PLOT
#             plot(Gray.(us_flow[:,:,1]), aspect_ratio=:equal, framestyle=:none, legend=false, yflip=false)
#             savefig(filenamepng)
#             next!(p)
#         end
#     end
# end
# cd(cwd)


import Base: isempty,
             rand,
             convert

export HPolyhedron,
       dim, σ, ∈,
       addconstraint!,
       constraints_list,
       tohrep, tovrep,
       convex_hull,
       cartesian_product,
       isempty,
       remove_redundant_constraints,
       remove_redundant_constraints!,
       constrained_dimensions,
       is_hyperplanar

"""
    HPolyhedron{N, VN<:AbstractVector{N}} <: AbstractPolyhedron{N}

Type that represents a convex polyhedron in H-representation, that is a finite intersection of half-spaces,
```math
P = \\bigcap_{i = 1}^m H_i,
```
where each ``H_i = \\{x \\in \\mathbb{R}^n : a_i^T x \\leq b_i \\}`` is a half-space,
``a_i \\in \\mathbb{R}^n`` is the normal vector of the ``i``-th half-space and ``b_i`` is the displacement.
The set ``P`` may or may not be bounded (see also [`HPolytope`](@ref), which assumes boundedness).

### Fields

- `constraints` -- vector of linear constraints
"""
struct HPolyhedron{N, VN<:AbstractVector{N}} <: AbstractPolyhedron{N}
    constraints::Vector{LinearConstraint{N, VN}}

    function HPolyhedron(constraints::Vector{LinearConstraint{N, VN}}) where {N,
                                                                              VN<:AbstractVector{N}}
        return new{N, VN}(constraints)
    end
end

isoperationtype(::Type{<:HPolyhedron}) = false
isconvextype(::Type{<:HPolyhedron}) = true

# constructor for an HPolyhedron with no constraints
function HPolyhedron{N, VN}() where {N, VN<:AbstractVector{N}}
    HPolyhedron(Vector{LinearConstraint{N, VN}}())
end

# constructor for an HPolyhedron with no constraints and given numeric type
function HPolyhedron{N}() where {N}
    HPolyhedron(Vector{LinearConstraint{N, Vector{N}}}())
end

# constructor for an HPolyhedron without explicit numeric type, defaults to Float64
function HPolyhedron()
    HPolyhedron{Float64}()
end

# constructor for an HPolyhedron with constraints of mixed type
function HPolyhedron(constraints::Vector{<:LinearConstraint})
    HPolyhedron(_normal_Vector(constraints))
end

# constructor from a simple H-representation
HPolyhedron(A::AbstractMatrix, b::AbstractVector) =
    HPolyhedron(constraints_list(A, b))

# convenience union type
const HPoly{N} = Union{HPolytope{N}, HPolyhedron{N}}

# --- LazySet interface functions ---

"""
    dim(P::HPoly)

Return the dimension of a polyhedron in H-representation.

### Input

- `P`  -- polyhedron in H-representation

### Output

The ambient dimension of the polyhedron in H-representation.
If it has no constraints, the result is ``-1``.
"""
function dim(P::HPoly)
    return length(P.constraints) == 0 ? -1 : length(P.constraints[1].a)
end

"""
    ρ(d::AbstractVector{M}, P::HPoly{N};
      solver=default_lp_solver(M, N)) where {M, N}

Evaluate the support function of a polyhedron (in H-representation) in a given
direction.

### Input

- `d`      -- direction
- `P`      -- polyhedron in H-representation
- `solver` -- (optional, default: `default_lp_solver(M, N)`) the backend used to
              solve the linear program

### Output

The support function of the polyhedron.
If a polytope is unbounded in the given direction, we throw an error.
If a polyhedron is unbounded in the given direction, the result is `Inf`.
"""
function ρ(d::AbstractVector{M}, P::HPoly{N};
           solver=default_lp_solver(M, N)) where {M, N}
    lp, unbounded = σ_helper(d, P, solver)
    if unbounded
        if P isa HPolytope
            error("the support function in direction $(d) is undefined " *
                  "because the polytope is unbounded")
        end
        return N(Inf)
    end
    return dot(d, lp.sol)
end

"""
    σ(d::AbstractVector{M}, P::HPoly{N};
      solver=default_lp_solver(M, N) where {M, N}

Return the support vector of a polyhedron (in H-representation) in a given
direction.

### Input

- `d`      -- direction
- `P`      -- polyhedron in H-representation
- `solver` -- (optional, default: `default_lp_solver(M, N)`) the backend used to
              solve the linear program

### Output

The support vector in the given direction.
"""
function σ(d::AbstractVector{M}, P::HPoly{N};
           solver=default_lp_solver(M, N)) where {M, N}
    lp, unbounded = σ_helper(d, P, solver)
    if unbounded
        if P isa HPolytope
            error("the support vector in direction $(d) is undefined because " *
                  "the polytope is unbounded")
        end
        # construct the solution from the solver's ray result
        if lp == nothing
            ray = d
        elseif has_lp_infeasibility_ray(lp.model)
            ray = lp.sol  # infeasibility ray is stored as the solution
        else
            error("LP solver did not return an infeasibility ray")
        end
        res = Vector{N}(undef, length(ray))
        @inbounds for i in 1:length(ray)
            if ray[i] == zero(N)
                res[i] = zero(N)
            elseif ray[i] > zero(N)
                res[i] = N(Inf)
            else
                res[i] = N(-Inf)
            end
        end
        return res
    else
        return lp.sol
    end
end

function σ_helper(d::AbstractVector, P::HPoly, solver)
    # let c = -d as a Vector since GLPK does not accept sparse vectors
    # (see #1011)
    c = to_negative_vector(d)

    (A, b) = tosimplehrep(P)
    if length(b) == 0
        unbounded = true
        lp = nothing
    else
        sense = '<'
        l = -Inf
        u = Inf
        lp = linprog(c, A, sense, b, l, u, solver)
        if is_lp_infeasible(lp.status, strict=true)
            error("the support vector is undefined because the polyhedron is " *
                  "empty")
        elseif is_lp_unbounded(lp.status)
            unbounded = true
        elseif is_lp_optimal(lp.status)
            unbounded = false
        else
            error("got unknown LP status $(lp.status)")
        end
    end
    return (lp, unbounded)
end

"""
    rand(::Type{HPolyhedron}; [N]::Type{<:Real}=Float64, [dim]::Int=2,
         [rng]::AbstractRNG=GLOBAL_RNG, [seed]::Union{Int, Nothing}=nothing)

Create a polyhedron.

### Input

- `HPolyhedron` -- type for dispatch
- `N`           -- (optional, default: `Float64`) numeric type
- `dim`         -- (optional, default: 2) dimension (is ignored)
- `rng`         -- (optional, default: `GLOBAL_RNG`) random number generator
- `seed`        -- (optional, default: `nothing`) seed for reseeding

### Output

A polyhedron.

### Algorithm

We first create a random polytope and then randomly remove some of the
constraints.
"""
function rand(::Type{HPolyhedron};
              N::Type{<:Real}=Float64,
              dim::Int=2,
              rng::AbstractRNG=GLOBAL_RNG,
              seed::Union{Int, Nothing}=nothing)
    rng = reseed(rng, seed)
    P = rand(HPolytope; N=N, dim=dim, rng=rng)
    constraints_P = constraints_list(P)
    constraints_Q = Vector{eltype(constraints_P)}()
    for i in 1:length(constraints_P)
        if rand(Bool)
            push!(constraints_Q, constraints_P[i])
        end
    end
    return HPolyhedron(constraints_Q)
end

# ===========================================
# HPolyhedron and HPolytope's shared methods
# ===========================================

"""
    addconstraint!(P::HPoly, constraint::LinearConstraint)

Add a linear constraint to a polyhedron in H-representation.

### Input

- `P`          -- polyhedron in H-representation
- `constraint` -- linear constraint to add

### Notes

It is left to the user to guarantee that the dimension of all linear constraints
is the same.
"""
function addconstraint!(P::HPoly, constraint::LinearConstraint)
    push!(P.constraints, constraint)
    return nothing
end

"""
    constraints_list(P::HPoly)

Return the list of constraints defining a polyhedron in H-representation.

### Input

- `P` -- polyhedron in H-representation

### Output

The list of constraints of the polyhedron.
"""
function constraints_list(P::HPoly)
    return P.constraints
end

"""
    tohrep(P::HPoly)

Return a constraint representation of the given polyhedron in constraint
representation (no-op).

### Input

- `P` -- polyhedron in constraint representation

### Output

The same polyhedron instance.
"""
function tohrep(P::HPoly)
    return P
end

"""
    normalize(P::HPoly{N}, p=N(2)) where {N}

Normalize a polyhedron in constraint representation.

### Input

- `P` -- polyhedron in constraint representation
- `p` -- (optional, default: `2`) norm

### Output

A new polyhedron in constraint representation whose normal directions ``a_i``
are normalized, i.e., such that ``‖a_i‖_p = 1`` holds.
"""
function normalize(P::HPoly{N}, p=N(2)) where {N}
    constraints = [normalize(hs, p) for hs in constraints_list(P)]
    T = basetype(P)
    return T(constraints)
end

"""
    remove_redundant_constraints(P::HPoly{N};
                                 backend=default_lp_solver(N)) where {N}

Remove the redundant constraints in a polyhedron in H-representation.

### Input

- `P`       -- polyhedron
- `backend` -- (optional, default: `default_lp_solver(N)`) the backend used to
               solve the linear program

### Output

A polyhedron equivalent to `P` but with no redundant constraints, or an empty
set if `P` is detected to be empty, which may happen if the constraints are
infeasible.

### Algorithm

See `remove_redundant_constraints!(::AbstractVector{<:LinearConstraint})` for
details.
"""
function remove_redundant_constraints(P::HPoly{N};
                                      backend=default_lp_solver(N)) where {N}
    Pred = copy(P)
    if remove_redundant_constraints!(Pred, backend=backend)
        return Pred
    else # the polyhedron P is empty
        return EmptySet{N}(dim(P))
    end
end

"""
    remove_redundant_constraints!(P::HPoly{N};
                                  backend=default_lp_solver(N)) where {N}

Remove the redundant constraints in a polyhedron in H-representation; the
polyhedron is updated in-place.

### Input

- `P`       -- polyhedron
- `backend` -- (optional, default: `default_lp_solver(N)`) the backend used to
               solve the linear program

### Output

`true` if the method was successful and the polyhedron `P` is modified by
removing its redundant constraints, and `false` if `P` is detected to be empty,
which may happen if the constraints are infeasible.

### Algorithm

See `remove_redundant_constraints!(::AbstractVector{<:LinearConstraint})` for
details.
"""
function remove_redundant_constraints!(P::HPoly{N};
                                       backend=default_lp_solver(N)) where {N}
    remove_redundant_constraints!(P.constraints, backend=backend)
end

"""
    translate(P::HPoly, v::AbstractVector; [share]::Bool=false)

Translate (i.e., shift) a polyhedron in constraint representation by a given
vector.

### Input

- `P`     -- polyhedron in constraint representation
- `v`     -- translation vector
- `share` -- (optional, default: `false`) flag for sharing unmodified parts of
             the original set representation

### Output

A translated polyhedron in constraint representation.

### Notes

The normal vectors of the constraints (vector `a` in `a⋅x ≤ b`) are shared with
the original constraints if `share == true`.

### Algorithm

We translate every constraint.
"""
function translate(P::HPoly, v::AbstractVector; share::Bool=false)
    @assert length(v) == dim(P) "cannot translate a $(dim(P))-dimensional " *
                                "set by a $(length(v))-dimensional vector"
    constraints = [translate(c, v; share=share) for c in constraints_list(P)]
    T = basetype(P)
    return T(constraints)
end

# ========================================================
# External methods that require Polyhedra.jl to be loaded
# ========================================================

"""
    tovrep(P::HPoly;
          [backend]=default_polyhedra_backend(P))

Transform a polyhedron in H-representation to a polytope in V-representation.

### Input

- `P`       -- polyhedron in constraint representation
- `backend` -- (optional, default: `default_polyhedra_backend(P)`) the
               backend for polyhedral computations

### Output

The `VPolytope` which is the vertex representation of the given polyhedron
in constraint representation.

### Notes

The conversion may not preserve the numeric type (e.g., with `N == Float32`)
depending on the backend.
For further information on the supported backends see [Polyhedra's
documentation](https://juliapolyhedra.github.io/).
"""
function tovrep(P::HPoly;
                backend=default_polyhedra_backend(P))
    require(:Polyhedra; fun_name="tovrep")
    P = polyhedron(P; backend=backend)
    return VPolytope(P)
end

"""
    isempty(P::HPoly{N}, witness::Bool=false;
            [use_polyhedra_interface]::Bool=false, [solver]=nothing,
            [backend]=nothing) where {N}

Determine whether a polyhedron is empty.

### Input

- `P`       -- polyhedron
- `witness` -- (optional, default: `false`) compute a witness if activated
- `use_polyhedra_interface` -- (optional, default: `false`) if `true`, we use
               the `Polyhedra` interface for the emptiness test
- `solver`  -- (optional, default: `nothing`) LP-solver backend, uses `default_lp_solver(N)`
               if it is not provided
- `backend` -- (optional, default: `nothing`) backend for polyhedral
               computations in `Polyhedra`, uses `default_polyhedra_backend(P)` if
               it is not provided

### Output

* If `witness` option is deactivated: `true` iff ``P = ∅``
* If `witness` option is activated:
  * `(true, [])` iff ``P = ∅``
  * `(false, v)` iff ``P ≠ ∅`` and ``v ∈ P``

### Notes

The default value of the `backend` is set internally and depends on whether the
`use_polyhedra_interface` option is set or not.
If the option is set, we use `default_polyhedra_backend(P)`.

Witness production is not supported if `use_polyhedra_interface` is `true`.

### Algorithm

The algorithm sets up a feasibility LP for the constraints of `P`.
If `use_polyhedra_interface` is `true`, we call `Polyhedra.isempty`.
Otherwise, we set up the LP internally.
"""
function isempty(P::HPoly{N},
                 witness::Bool=false;
                 use_polyhedra_interface::Bool=false,
                 solver=nothing,
                 backend=nothing) where {N}
    if length(constraints_list(P)) < 2
        # catch corner case because of problems in LP solver for Rationals
        return witness ? (false, an_element(P)) : false
    end
    if use_polyhedra_interface
        require(:Polyhedra; fun_name="isempty", explanation="with the active " *
            "option `use_polyhedra_interface`")

        if backend == nothing
            backend = default_polyhedra_backend(P)
        end

        if isnothing(solver)
            result = Polyhedra.isempty(polyhedron(P; backend=backend))
        else
            result = Polyhedra.isempty(polyhedron(P; backend=backend), solver)
        end

        if result
            return witness ? (true, N[]) : true
        elseif witness
            error("witness production is not supported yet")
        else
            return false
        end
    else
        A, b = tosimplehrep(P)
        lbounds, ubounds = -Inf, Inf
        sense = '<'
        obj = zeros(N, size(A, 2))
        if isnothing(solver)
            solver = default_lp_solver(N)
        end
        lp = linprog(obj, A, sense, b, lbounds, ubounds, solver)
        if is_lp_optimal(lp.status)
            return witness ? (false, lp.sol) : false
        elseif is_lp_infeasible(lp.status)
            return witness ? (true, N[]) : true
        end
        error("LP returned status $(lp.status) unexpectedly")
    end
end

# ==================================
# Methods that use Polyhedra.jl
# ==================================

function load_polyhedra_hpolyhedron() # function to be loaded by Requires
return quote
# see the interface file init_Polyhedra.jl for the imports

function convert(::Type{HPolyhedron}, P::HRep{N}) where {N}
    VN = Polyhedra.hvectortype(P)
    constraints = Vector{LinearConstraint{N, VN}}()
    for hi in Polyhedra.allhalfspaces(P)
        a, b = hi.a, hi.β
        if isapproxzero(norm(a))
            continue
        end
        push!(constraints, HalfSpace(a, b))
    end
    return HPolyhedron(constraints)
end

"""
     HPolyhedron(P::HRep{N}) where {N}

Return a polyhedron in H-representation given a `HRep` polyhedron
from `Polyhedra.jl`.

### Input

- `P` -- `HRep` polyhedron

### Output

An `HPolyhedron`.
"""
function HPolyhedron(P::HRep{N}) where {N}
    convert(HPolyhedron, P)
end

"""
    polyhedron(P::HPoly; [backend]=default_polyhedra_backend(P))

Return an `HRep` polyhedron from `Polyhedra.jl` given a polytope in
H-representation.

### Input

- `P`       -- polytope
- `backend` -- (optional, default: call `default_polyhedra_backend(P)`)
                the polyhedral computations backend

### Output

An `HRep` polyhedron.

### Notes

For further information on the supported backends see
[Polyhedra's documentation](https://juliapolyhedra.github.io/).
"""
function polyhedron(P::HPoly;
                    backend=default_polyhedra_backend(P))
    A, b = tosimplehrep(P)
    return Polyhedra.polyhedron(Polyhedra.hrep(A, b), backend)
end

function triangulate(X::LazySet)

    dim(X) == 3 || throw(ArgumentError("the dimension of the set should be three, got $(dim(X))"))

    poly = polyhedron(convert(HPolyhedron, X))
    mes = Mesh(poly)
    coords = Polyhedra.GeometryBasics.coordinates(mes)
    connec = Polyhedra.GeometryBasics.faces(mes)

    ntriangles = length(connec)
    npoints = 3*ntriangles
    points = Matrix{Float32}(undef, 3, npoints)

    for i in 1:npoints
        points[:, i] .= coords[i].data
    end

    connec_tup = getfield.(connec, :data)

    return points, connec_tup
end

end # quote
end # function load_polyhedra_hpolyhedron()

"""
    _isbounded_stiemke(P::HPolyhedron{N}; solver=LazySets.default_lp_solver(N),
                       check_nonempty::Bool=true) where {N}

Determine whether a polyhedron is bounded using Stiemke's theorem of alternatives.

### Input

- `P`       -- polyhedron
- `backend` -- (optional, default: `default_lp_solver(N)`) the backend used
               to solve the linear program
- `check_nonempty` -- (optional, default: `true`) if `true`, check the
                      precondition to this algorithm that `P` is non-empty

### Output

`true` iff the polyhedron is bounded

### Notes

The algorithm internally calls `isempty` to check whether the polyhedron is empty.
This computation can be avoided using the `check_nonempty` flag.

### Algorithm

The algorithm is based on Stiemke's theorem of alternatives, see e.g. [1].

Let the polyhedron ``P`` be given in constraint form ``Ax ≤ b``. We assume that
the polyhedron is not empty.

Proposition 1. If ``\\ker(A)≠\\{0\\}``, then ``P`` is unbounded.

Proposition 2. Assume that ``ker(A)={0}`` and ``P`` is non-empty.
Then ``P`` is bounded if and only if the following linear
program admits a feasible solution: ``\\min∥y∥_1`` subject to ``A^Ty=0`` and ``y≥1``.

[1] Mangasarian, Olvi L. *Nonlinear programming.*
    Society for Industrial and Applied Mathematics, 1994.
"""
function _isbounded_stiemke(P::HPolyhedron{N}; solver=LazySets.default_lp_solver(N),
                            check_nonempty::Bool=true) where {N}
    if check_nonempty && isempty(P)
        return true
    end

    A, b = tosimplehrep(P)
    m, n = size(A)

    if !isempty(nullspace(A))
        return false
    end

    At = copy(transpose(A))
    c = ones(N, m)
    lp = linprog(c, At, '=', zeros(n), one(N), Inf, solver)
    return is_lp_optimal(lp.status)
end

function is_hyperplanar(P::HPolyhedron)
    clist = P.constraints
    m = length(clist)

    # check that the number of constraints is fine
    if m > 2
        # try to remove redundant constraints
        clist = remove_redundant_constraints(clist)
        m = length(clist)
    end
    if m != 2
        return false
    end

    # check that the two half-spaces are complementary
    return @inbounds iscomplement(clist[1], clist[2])
end

# ============================================
# Functionality that requires Symbolics
# ============================================
function load_symbolics_hpolyhedron()

return quote

"""
    HPolyhedron(expr::Vector{<:Num}, vars=_get_variables(expr); N::Type{<:Real}=Float64)

Return the polyhedron in half-space representation given by a list of symbolic expressions.

### Input

- `expr` -- vector of symbolic expressions that describes each half-space
- `vars` -- (optional, default: `_get_variables(expr)`), if an array of variables is given,
            use those as the ambient variables in the set with respect to which derivations
            take place; otherwise, use only the variables which appear in the given
            expression (but be careful because the order may be incorrect; it is advicsed to always
            pass `vars` explicitly)
- `N`    -- (optional, default: `Float64`) the numeric type of the returned half-space

### Output

An `HPolyhedron`.

### Examples

```jldoctest
julia> using Symbolics

julia> vars = @variables x y
2-element Vector{Num}:
 x
 y

julia> HPolyhedron([x + y <= 1, x + y >= -1], vars)
HPolyhedron{Float64, Vector{Float64}}(HalfSpace{Float64, Vector{Float64}}[HalfSpace{Float64, Vector{Float64}}([1.0, 1.0], 1.0), HalfSpace{Float64, Vector{Float64}}([-1.0, -1.0], 1.0)])

julia> X = HPolyhedron([x == 0, y <= 0], vars)
HPolyhedron{Float64, Vector{Float64}}(HalfSpace{Float64, Vector{Float64}}[HalfSpace{Float64, Vector{Float64}}([1.0, 0.0], -0.0), HalfSpace{Float64, Vector{Float64}}([-1.0, -0.0], 0.0), HalfSpace{Float64, Vector{Float64}}([0.0, 1.0], -0.0)])
```
"""
function HPolyhedron(expr::Vector{<:Num}, vars::AbstractVector{Num}; N::Type{<:Real}=Float64)
    clist = Vector{HalfSpace{N, Vector{N}}}()
    sizehint!(clist, length(expr))
    got_hyperplane = false
    got_halfspace = false
    zeroed_vars = Dict(v => zero(N) for v in vars)
    vars_list = collect(vars)
    for ex in expr
        exval = Symbolics.value(ex)
        got_hyperplane, sexpr = _is_hyperplane(exval)
        if !got_hyperplane
            got_halfspace, sexpr = _is_halfspace(exval)
            if !got_halfspace
                throw(ArgumentError("expected an expression describing either " *
                    "a half-space of a hyperplane, got $expr"))
            end
        end

        coeffs = [N(α.val) for α in gradient(sexpr, vars_list)]
        β = -N(Symbolics.substitute(sexpr, zeroed_vars))

        push!(clist, HalfSpace(coeffs, β))
        if got_hyperplane
            push!(clist, HalfSpace(-coeffs, -β))
        end
    end
    return HPolyhedron(clist)
end

HPolyhedron(expr::Vector{<:Num}; N::Type{<:Real}=Float64) = HPolyhedron(expr, _get_variables(expr); N=N)
HPolyhedron(expr::Vector{<:Num}, vars; N::Type{<:Real}=Float64) = HPolyhedron(expr, _vec(vars); N=N)

end end  # quote / load_symbolics_hpolyhedron()


import SparseArrays: SparseVector, sparsevec
import LinearAlgebra: cholesky, Matrix

struct BasisFunction{PT,RT}
    center::PT
    coefficients::SparseVector{RT, Int}
end

function BasisFunction(center, coefficients)
    return BasisFunction{typeof(center), eltype(coefficients)}(center, coefficients)
end

function real_type(::Type{BasisFunction{PT, RT}})  where {PT, RT}
    return RT
end

function center(bf::BasisFunction)
    return bf.center
end

function coefficients(bf::BasisFunction)
    return bf.coefficients
end

function create_basis_vector(id_arrays, weights, N)
    @assert size(id_arrays) == size(weights)
    out_weights = vcat([weights[k] * one.(id_arrays[k]) for k = 1 : length(weights)]...)
    return sparsevec(vcat(id_arrays...), out_weights, N)
end

# Input are the coarsest-scale domains
function compute_basis_functions(domains::AbstractVector{Domain{PT}}, N = maximum(id.(gather_descendants(domains)))) where PT<:AbstractVector{<:Real}
    centers = Vector{Vector{PT}}(undef,1)
    centers[1] = Vector{PT}(undef, 0)
    basis_vectors = Vector{Vector{SparseVector{eltype(PT), Int}}}(undef,1)
    basis_vectors[1] = Vector{SparseVector{eltype(PT), Int}}(undef, 0)
    # recursive functions that computes the basis functions orthogonal in ℓ^2 inner product. Here, k is the level of domain
    function recurse(dm, k)
        # If the domain is elementrary, we return an array containint its index.
        n_ch = length(children(dm))
        if iselementary(dm)
            return [id(dm)]
        else
            # Extend length of output arrays if first domain on new level
            if k + 1 > length(basis_vectors)
                # make sure we didn't skip a level
                @assert length(basis_vectors) == length(centers) == k
                push!(centers, Vector{PT}(undef, 0))
                push!(basis_vectors, Vector{SparseVector{eltype(PT), Int}}(undef, 0))
            end

            # Compute coefficients of linear combination
            M = zeros(eltype(PT), n_ch + 1, n_ch + 1)
            M[1, 1] = sum(weight.(children(dm)))
            # setting the inner products.
            for i = 1 : n_ch;
                M[1, i + 1] = M[i + 1, 1] = M[i + 1, i + 1] = weight(children(dm)[i])
            end
            # Adding 1 to the last entry of M to ensure positivity#
            # This should not affect the result, since the affected columns are dropped
            M[end, end] += 1
            # Compute QR factorization that will provide coefficients
            R = inv(Matrix(cholesky(M).L'))

            # apply the function to the children to obtain the children arrays
            id_arrays = map(ch -> recurse(ch, k + 1), children(dm))
            id_arrays = vcat([vcat(id_arrays...)], id_arrays)

            # If on first level, return normalized average as basis function

            if k == 1
                push!(basis_vectors[k], create_basis_vector(id_arrays, R[:, 1], N))
                push!(centers[k], center(dm))
            end

            # add the remaining basis functions
            for ch = 2 : n_ch
                push!(basis_vectors[k + 1], create_basis_vector(id_arrays, R[:, ch], N))
                push!(centers[k + 1], center(dm))
            end
            return id_arrays[1]
        end
    end
    for dm in domains
        recurse(dm, 1)
    end

    # combine basis_vectors and centers to basis functions
    # cannot broadcast directly, because we are dealing with arrays of arrays of basis functions due to the different scales.
    out = Vector{Vector{BasisFunction{eltype(centers[1]), eltype(basis_vectors[1][1])}}}(undef, length(basis_vectors))
    for k = 1 : length(basis_vectors)
        out[k] = BasisFunction.(centers[k], basis_vectors[k])
    end

    return out
end


module DA_DCSTests

using Test
using LightPropagation: g2_DA_semiinf_CW, g2_DA_Nlay_cylinder_CW
using LightPropagation: DAsemiinf_DCS, Nlayer_cylinder_DCS, besselroots

τ = 10 .^(range(-10,stop=0,length=250))
ρ = 1.0; μa = 0.1; μsp = 10.0; n_med = 1.0; n_ext = 1.0
β = 1.0; λ = 700.0; z = 0.0
BFi = 2.0e-8

data = DAsemiinf_DCS(ρ = ρ, μa = μa, μsp = μsp, n_med = n_med, n_ext = n_ext, β = β, λ = λ, BFi = BFi, z = z)

@test g2_DA_semiinf_CW(τ, data) ≈ g2_DA_semiinf_CW(τ, ρ, μa, μsp; BFi = BFi, β = β, n_ext = n_ext, n_med = n_med, z = z, λ = λ)

si = g2_DA_semiinf_CW(τ, data)

### test layered against semi-inf

μa = [0.1, 0.1]; μsp = [10.0, 10.0]; n_med = [1.0, 1.0]; n_ext = 1.0
BFi = [2.0e-8, 2.0e-8]; l = [1.0, 10.0]; a = 25.0

data = Nlayer_cylinder_DCS(ρ = ρ, μa = μa, μsp = μsp, n_med = n_med, n_ext = n_ext, β = β, λ = λ, BFi = BFi, z = z, a = a, l = l, bessels = besselroots[1:2000])

@test g2_DA_Nlay_cylinder_CW(τ, data) ≈ g2_DA_Nlay_cylinder_CW(τ, ρ, μa, μsp; BFi = BFi, β = β, n_ext = n_ext, n_med = n_med, l = l, a = a, z = z, λ = λ, bessels = besselroots[1:2000])

layered = g2_DA_Nlay_cylinder_CW(τ, data)

@test si ≈ layered

end


using QCMaterial
using PyCall
using LinearAlgebra

scipy_opt = pyimport("scipy.optimize")
struct Hamiltonian
    ham
    n_sites
end

hfstate(n_qubit, n_electron) = parse(Int, repeat("0", n_qubit-n_electron) * repeat("1", n_electron), base=2)

function generate_ham(nsite)
    # ## 2-site Hubbard modelのハミルトニアンを定義する。
    #
    # $H=-t \sum_{\langle i, j>\sigma=(\uparrow, \downarrow)} \sum_{i, \sigma}\left(a_{i, \sigma}^{\dagger} a_{j, \sigma}+a_{j, \sigma}^{\dagger} a_{i, \sigma}\right)+U \sum_{i} n_{i, \uparrow} n_{i, \downarrow}-\mu \sum_{i} \sum_{\sigma=(\uparrow, \downarrow)} n_{i, \sigma}$
    #
    # $1_\uparrow$->0,$1_\downarrow$->1, $2_\uparrow$->2, $2_\downarrow$->3
    #
    # クーロン斥力(U=2): 2 * $n_{1,\uparrow}$ $n_{1,\downarrow}$ + 2 * $n_{2,\uparrow}$ $n_{2,\downarrow}$
    #                = 2 * $c^\dagger_0$ $c_0$ $c^\dagger_1$ $c_1$ + 2 * $c^\dagger_2$ $c_2$ $c^\dagger_3$ $c_3$
    #
    # ホッピング項(t=-0.1): -0.1 * $c^\dagger_{1,\uparrow}$ $c_{2,\uparrow}$ -0.1 * $c^\dagger_{2,\uparrow}$ $c_{1,\uparrow}$
    #                    -0.1 * $c^\dagger_{1,\downarrow}$ $c_{2,\downarrow}$ -0.1 * $c^\dagger_{2,\downarrow}$ $c_{1,\downarrow}$
    #   = -0.1 * $c^\dagger_0$ $c_2$ + -0.1 * $c^\dagger_2$ $c_0$ + -0.1 * $c^\dagger_1$ $c_3$ + -0.1 * $c_3^\dagger$ $c_1$
    #
    # ケミカルポテンシャル項(μ=U/2=1):-1 * $n_{1,\uparrow}$ - 1 * $n_{1,\downarrow}$ - 1 * $n_{2,\uparrow}$ - 1 * $n_{2,\downarrow}$
    #                = - 1 * $c^\dagger_0$ $c_0$ - 1 * $c^\dagger_1$ $c_1$ - 1 * $c^\dagger_2$ $c_2$ - 1 * $c^\dagger_3$ $c_3$

    #ハバードモデルハミルトニアンを定義し、対角化まで行う
    n_qubit = 2*nsite #量子ビットの数
    U = 10.0
    t = -1.0

    #ハーフフィリングを仮定(電子数 = サイトの数)するとケミカルポテンシャルはμ=U/2.マイナスがつくので以下。
    μ = U/2

    ham = FermionOperator()
    #斥力項
    for i in 1:nsite
        #up_index,down_indexの定義は、QC_materialを参照。
        up = up_index(i)
        down = down_index(i)
        ham += FermionOperator("$(up)^ $(down)^ $(up) $(down)", -U) #左側に生成演算子。右際に消滅演算子をもっていく過程で半交換関係が1回でマイナスをつける。
    end
    #ホッピング項
    for i in 1:nsite-1
        ham += FermionOperator("$(up_index(i+1))^ $(up_index(i))", t)
        ham += FermionOperator("$(up_index(i))^ $(up_index(i+1))", t)
        ham += FermionOperator("$(down_index(i+1))^ $(down_index(i))", t)
        ham += FermionOperator("$(down_index(i))^ $(down_index(i+1))", t)
    end

    #ケミカルポテンシャルの項
    for i in 1:nsite
        up = up_index(i)
        down = down_index(i)
        ham += FermionOperator("$(up)^  $(up) ", -μ)
        ham += FermionOperator("$(down)^ $(down)", -μ)
    end

    Hamiltonian(ham, nsite)
end

"""
Evaluate energy
"""
function eval_energy(circuit, ham_jw, theta_list, n_qubit)
    state = create_hf_state(n_qubit, n_electron) #|0000> を準備
    update_circuit_param!(circuit, theta_list) #量子回路にパラメータをセット
    update_quantum_state!(circuit, state) #量子回路を状態に作用
    get_expectation_value(ham_jw, state) #ハミルトニアンの期待値
end

function construct_circuit(hamiltonian)
    ham = hamiltonian.ham
    n_qubit = 2*hamiltonian.n_sites

    @assert mod(n_electron, 2) == 0
    @assert n_electron <= hamiltonian.n_sites

    #JW変換
    ham_jw = jordan_wigner(ham)

    # Prepare a circuit
    circuit = uccgsd(n_qubit, orbital_rot=true)
    circuit, ham_jw
end

function solve(hamiltonian, n_electron;theta_init=nothing, comm=nothing)
    if comm === nothing
        rank = 0
    else
        rank = MPI.Comm_rank(comm)
    end

    ham = hamiltonian.ham
    n_qubit = 2*hamiltonian.n_sites

    println(ham)　#ハバードモデルハミルトニアンの表示。
    @assert mod(n_electron, 2) == 0
    @assert n_electron <= hamiltonian.n_sites

    # Compute exact ground-state energy
    sparse_mat = get_number_preserving_sparse_operator(ham, n_qubit, n_electron)　#行列の取得
    enes_ed = eigvals(sparse_mat.toarray())　#対角化を行う
    EigVal_min = minimum(enes_ed)
    if rank == 0
        println("Exact ground-state energy = $(EigVal_min)")
    end

    # Construct circuit
    circuit, ham_jw = construct_circuit(hamiltonian)
    if rank == 0
        println("Number of Qubits:", n_qubit)
        println("Number of Electrons:", n_electron)
    end

    # Define a cost function
    function cost(theta_list)
        eval_energy(circuit, ham_jw, theta_list, n_qubit)
    end

    # Define the gradient of the cost function
    function grad_cost(theta_list)
        t1 = time_ns()
        if comm === nothing
            first_idx, size = 1, length(theta_list)
        else
            first_idx, size = distribute(length(theta_list), MPI.Comm_size(comm), MPI.Comm_rank(comm))
        end
        last_idx = first_idx + size - 1
        res = numerical_grad(cost, theta_list, first_idx=first_idx, last_idx=last_idx)
        if comm !== nothing
            res = MPI.Allreduce(res, MPI.SUM, comm)
        end
        t2 = time_ns()
        if rank == 0
            println("g: ", (t2-t1)*1e-9)
        end
        res
    end

    if theta_init === nothing
        theta_init = rand(size(circuit.theta_offsets)[1])
        if comm !== nothing
            # Make sure all processes use the same initial values
            MPI.Bcast!(theta_init, 0, comm)
        end
    end
    cost_history = Float64[] #コスト関数の箱
    init_theta_list = theta_init
    push!(cost_history, cost(init_theta_list))

    method = "BFGS"
    options = Dict("disp" => true, "maxiter" => 200, "gtol" => 1e-5)
    function callback(x)
        push!(cost_history, cost(x))
        if rank == 0
            println("iter ", length(cost_history), " ", cost_history[end])
        end
    end
    opt = scipy_opt.minimize(cost, init_theta_list, method=method, callback=callback, jac=grad_cost,
        options=options)

    return cost_history, circuit, EigVal_min, opt
end


using ITensors
using Test
import Combinatorics

digits(::Type{T},i,j,k) where {T} = T(i*10^2+j*10+k)

@testset "ITensor $T Contractions" for T ∈ (Float64,ComplexF64)
  mi,mj,mk,ml,mα = 2,3,4,5,6,7
  i = Index(mi,"i")
  j = Index(mj,"j")
  k = Index(mk,"k")
  l = Index(ml,"l")
  α = Index(mα,"alpha")
  @testset "Test contract ITensors" begin
    A = randomITensor(T)
    B = randomITensor(T)
    Ai = randomITensor(T,i)
    Bi = randomITensor(T,i)
    Aj = randomITensor(T,j)
    Aij = randomITensor(T,i,j)
    Bij = randomITensor(T,i,j)
    Aik = randomITensor(T,i,k)
    Ajk = randomITensor(T,j,k)
    Ajl = randomITensor(T,j,l)
    Akl = randomITensor(T,k,l)
    Aijk = randomITensor(T,i,j,k)
    Ajkl = randomITensor(T,j,k,l)
    Aikl = randomITensor(T,i,k,l)
    Aklα = randomITensor(T,k,l,α)
    Aijkl = randomITensor(T,i,j,k,l)
    @testset "Test contract ITensor (Scalar*Scalar -> Scalar)" begin
      C = A*B
      @test scalar(C)≈scalar(A)*scalar(B)
    end
    @testset "Test contract ITensor (Scalar*Vector -> Vector)" begin
      C = A*Ai
      @test array(C)≈scalar(A)*array(Ai)
    end
    @testset "Test contract ITensor (Vector*Scalar -> Vector)" begin
      C = Aj*A
      @test array(C)≈scalar(A)*array(Aj)
    end
    @testset "Test contract ITensors (Vectorᵀ*Vector -> Scalar)" begin
      C = Ai*Bi
      CArray = transpose(array(Ai))*array(Bi)
      @test CArray≈scalar(C)
    end
    @testset "Test contract ITensors (Vector*Vectorᵀ -> Matrix)" begin
      C = Ai*Aj
      for ii ∈ 1:dim(i), jj ∈ 1:dim(j)
        @test C[i(ii),j(jj)] ≈ Ai[i(ii)]*Aj[j(jj)]
      end
    end
    @testset "Test contract ITensors (Matrix*Scalar -> Matrix)" begin
      Aij = permute(Aij,i,j)
      C = Aij*A
      @test array(permute(C,i,j))≈scalar(A)*array(Aij)
    end
    @testset "Test contract ITensors (Matrix*Vector -> Vector)" begin
      Aij = permute(Aij,i,j)
      C = Aij*Aj
      CArray = array(permute(Aij,i,j))*array(Aj)
      @test CArray≈array(C)
    end
    @testset "Test contract ITensors (Matrixᵀ*Vector -> Vector)" begin
      Aij = permute(Aij,j,i)
      C = Aij*Aj
      CArray = transpose(array(Aij))*array(Aj)
      @test CArray≈array(C)
    end
    @testset "Test contract ITensors (Vector*Matrix -> Vector)" begin
      Aij = permute(Aij,i,j)
      C = Ai*Aij
      CArray = transpose(transpose(array(Ai))*array(Aij))
      @test CArray≈array(C)
    end
    @testset "Test contract ITensors (Vector*Matrixᵀ -> Vector)" begin
      Aij = permute(Aij,j,i)
      C = Ai*Aij
      CArray = transpose(transpose(array(Ai))*transpose(array(Aij)))
      @test CArray≈array(C)
    end
    @testset "Test contract ITensors (Matrix*Matrix -> Scalar)" begin
      Aij = permute(Aij,i,j)
      Bij = permute(Bij,i,j)
      C = Aij*Bij
      CArray = LinearAlgebra.tr(array(Aij)*transpose(array(Bij)))
      @test CArray≈scalar(C)
    end
    @testset "Test contract ITensors (Matrix*Matrix -> Matrix)" begin
      Aij = permute(Aij,i,j)
      Ajk = permute(Ajk,j,k)
      C = Aij*Ajk
      CArray = array(Aij)*array(Ajk)
      @test CArray≈array(C)
    end
    @testset "Test contract ITensors (Matrixᵀ*Matrix -> Matrix)" begin
      Aij = permute(Aij,j,i)
      Ajk = permute(Ajk,j,k)
      C = Aij*Ajk
      CArray = transpose(array(Aij))*array(Ajk)
      @test CArray≈array(C)
    end
    @testset "Test contract ITensors (Matrix*Matrixᵀ -> Matrix)" begin
      Aij = permute(Aij,i,j)
      Ajk = permute(Ajk,k,j)
      C = Aij*Ajk
      CArray = array(Aij)*transpose(array(Ajk))
      @test CArray≈array(C)
    end
    @testset "Test contract ITensors (Matrixᵀ*Matrixᵀ -> Matrix)" begin
      Aij = permute(Aij,j,i)
      Ajk = permute(Ajk,k,j)
      C = Aij*Ajk
      CArray = transpose(array(Aij))*transpose(array(Ajk))
      @test CArray≈array(C)
    end
    @testset "Test contract ITensors (Matrix⊗Matrix -> 4-tensor)" begin
      C = Aij*Akl
      for ii ∈ 1:dim(i), jj ∈ 1:dim(j), kk ∈ 1:dim(k), ll ∈ 1:dim(l)
        @test C[i(ii),j(jj),k(kk),l(ll)] ≈ Aij[i(ii),j(jj)]*Akl[k(kk),l(ll)]
      end
    end
    @testset "Test contract ITensors (3-Tensor*Scalar -> 3-Tensor)" begin
      Aijk = permute(Aijk,i,j,k)
      C = Aijk*A
      @test array(permute(C,i,j,k)) ≈ scalar(A)*array(Aijk)
    end
    @testset "Test contract ITensors (3-Tensor*Vector -> Matrix)" begin
      Aijk = permute(Aijk,i,j,k)
      C = Aijk*Ai
      CArray = reshape(reshape(array(permute(Aijk,j,k,i)),dim(j)*dim(k),dim(i))*array(Ai),dim(j),dim(k))
      @test CArray≈array(permute(C,j,k))
    end
    @testset "Test contract ITensors (Vector*3-Tensor -> Matrix)" begin
      Aijk = permute(Aijk,i,j,k)
      C = Aj*Aijk
      CArray = reshape(transpose(array(Aj))*reshape(array(permute(Aijk,j,i,k)),dim(j),dim(i)*dim(k)),dim(i),dim(k))
      @test CArray≈array(permute(C,i,k))
    end
    @testset "Test contract ITensors (3-Tensor*Matrix -> Vector)" begin
      Aijk = permute(Aijk,i,j,k)
      Aik = permute(Aik,i,k)
      C = Aijk*Aik
      CArray = reshape(array(permute(Aijk,j,i,k)),dim(j),dim(i)*dim(k))*vec(array(Aik))
      @test CArray≈array(C)
    end
    @testset "Test contract ITensors (3-Tensor*Matrix -> 3-Tensor)" begin
      Aijk = permute(Aijk,i,j,k)
      Ajl = permute(Ajl,j,l)
      C = Aijk*Ajl
      CArray = reshape(reshape(array(permute(Aijk,i,k,j)),dim(i)*dim(k),dim(j))*array(Ajl),dim(i),dim(k),dim(l))
      @test CArray≈array(permute(C,i,k,l))
    end
    @testset "Test contract ITensors (Matrix*3-Tensor -> 3-Tensor)" begin
      Aijk = permute(Aijk,i,j,k)
      Akl = permute(Akl,k,l)
      C = Akl*Aijk
      CArray = reshape(array(permute(Akl,l,k))*reshape(array(permute(Aijk,k,i,j)),dim(k),dim(i)*dim(j)),dim(l),dim(i),dim(j))
      @test CArray≈array(permute(C,l,i,j))
    end
    @testset "Test contract ITensors (3-Tensor*3-Tensor -> 3-Tensor)" begin
      Aijk = permute(Aijk,i,j,k)
      Ajkl = permute(Ajkl,j,k,l)
      C = Aijk*Ajkl
      CArray = reshape(array(permute(Aijk,i,j,k)),dim(i),dim(j)*dim(k))*reshape(array(permute(Ajkl,j,k,l)),dim(j)*dim(k),dim(l))
      @test CArray≈array(permute(C,i,l))
    end
    @testset "Test contract ITensors (3-Tensor*3-Tensor -> 3-Tensor)" begin
      for inds_ijk ∈ Combinatorics.permutations([i,j,k]), inds_jkl ∈ Combinatorics.permutations([j,k,l])
        Aijk = permute(Aijk,inds_ijk...)
        Ajkl = permute(Ajkl,inds_jkl...)
        C = Ajkl*Aijk
        CArray = reshape(array(permute(Ajkl,l,j,k)),dim(l),dim(j)*dim(k))*reshape(array(permute(Aijk,j,k,i)),dim(j)*dim(k),dim(i))
        @test CArray≈array(permute(C,l,i))
      end
    end
    @testset "Test contract ITensors (4-Tensor*3-Tensor -> 1-Tensor)" begin
      for inds_ijkl ∈ Combinatorics.permutations([i,j,k,l]), inds_jkl ∈ Combinatorics.permutations([j,k,l])
        Aijkl = permute(Aijkl,inds_ijkl...)
        Ajkl = permute(Ajkl,inds_jkl...)
        C = Ajkl*Aijkl
        CArray = reshape(array(permute(Ajkl,j,k,l)),1,dim(j)*dim(k)*dim(l))*reshape(array(permute(Aijkl,j,k,l,i)),dim(j)*dim(k)*dim(l),dim(i))
        @test vec(CArray)≈array(permute(C,i))
      end
    end
    @testset "Test contract ITensors (4-Tensor*3-Tensor -> 3-Tensor)" begin
      for inds_ijkl ∈ Combinatorics.permutations([i,j,k,l]), inds_klα ∈ Combinatorics.permutations([k,l,α])
        Aijkl = permute(Aijkl,inds_ijkl...)
        Aklα = permute(Aklα,inds_klα...)
        C = Aklα*Aijkl
        CArray = reshape(reshape(array(permute(Aklα,α,k,l)),dim(α),dim(k)*dim(l))*reshape(array(permute(Aijkl,k,l,i,j)),dim(k)*dim(l),dim(i)*dim(j)),dim(α),dim(i),dim(j))
        @test CArray≈array(permute(C,α,i,j))
      end
    end
  end # End contraction testset
end

@testset "Contraction conversions" begin

  @testset "Real scalar * Complex ITensor" begin
    i = Index(2,"i")
    j = Index(2,"j")
    x = rand(Float64)
    A = randomITensor(ComplexF64,i,j)
    B = x*A
    for ii ∈ dim(i), jj ∈ dim(j)
      @test B[i(ii),j(jj)] == x*A[i(ii),j(jj)]
    end
  end
  @testset "Complex scalar * Real ITensor" begin
    i = Index(2,"i")
    j = Index(2,"j")
    x = rand(ComplexF64)
    A = randomITensor(Float64,i,j)
    B = x*A
    for ii ∈ dim(i), jj ∈ dim(j)
      @test B[i(ii),j(jj)] == x*A[i(ii),j(jj)]
    end
  end
  @testset "Real ITensor * Complex ITensor" begin
    i = Index(2,"i")
    j = Index(2,"j")
    k = Index(2,"k")
    A = randomITensor(Float64,i,j)
    B = randomITensor(ComplexF64,j,k)
    C = A*B
    @test array(permute(C,i,k)) ≈ array(A)*array(B)
  end
  @testset "Complex ITensor * Real ITensor" begin
    i = Index(2,"i")
    j = Index(2,"j")
    k = Index(2,"k")
    A = randomITensor(ComplexF64,i,j)
    B = randomITensor(Float64,j,k)
    C = A*B
    @test array(permute(C,i,k)) ≈ array(A)*array(B)
  end

  @testset "Outer Product Real ITensor * Complex ITensor" begin
    i = Index(2, "i")
    j = Index(2, "j")
    A = randomITensor(Float64, i)
    B = randomITensor(ComplexF64, j)
    C = A*B
    @test array(permute(C, i, j)) ≈ kron(array(A), transpose(array(B)))
  end

  @testset "Outer Product: Complex ITensor * Real ITensor" begin
    i = Index(2, "i")
    j = Index(2, "j")
    A = randomITensor(ComplexF64, i)
    B = randomITensor(Float64, j)
    C = A*B
    @test array(permute(C, i, j)) ≈ kron(array(A), transpose(array(B)))
  end

end

nothing


module hetLK
struct VIterable
    upp::Array{Int}
    low::Array{Int}
    function VIterable(upp::Array{Int},low=fill(0,size(upp)))
        if length(upp)!=length(low) error("upp and low has to have the same length") end
        new(upp,low)
    end
end

function viterate(upp::Vector{Int},low=fill(0,size(upp)))
    Iterators.product(range.(low,upp,step=1)...)
end

# Vector Iterator
function Base.iterate(iter::VIterable,state=(iter.low,1))
    el, count = state
    if count>prod(iter.upp.-iter.low.+1) return nothing end
    nextel=copy(el)
    nextel[1]+=1
    j=1
    while j<length(nextel)
        if nextel[j]>iter.upp[j]
            nextel[j+1]+=1
            nextel[j]=iter.low[j]
        else
            break
        end
        j+=1
    end
    return (el, (nextel,count+1))
end
Base.length(iter::VIterable)=prod(iter.upp.-iter.low.+1)
Base.eltype(iter::VIterable)=Array{Int}

using Distributions
using LinearAlgebra
using StatsFuns
function LKlatent(i::Array{Int,1},n::Array{Int,1},lambda::Array{Float64,1},Rmat::Array{Float64,2},gamma::Float64,p::Array{Float64,1},memo::Dict{Pair{Vector{Int},Vector{Float64}},Float64},linememo::Dict{Pair{Vector{Int},Vector{Int}},Float64})
    if all(p.==1)
        return LK1(i,n,lambda,Rmat,gamma,memo,linememo)
    else
        ll=0.0
        p=fill(0.0,size(i)).+p
        for lat_i in VIterable(n,i)
            ll+=LK1(lat_i,n,lambda,Rmat,gamma,memo,linememo)+sum(logpdf.(Binomial.(lat_i,p),i))
            return ll
        end
    end
end

function LKlatent(i::SubArray{Int,1},n::SubArray{Int,1},lambda::Array{Float64,1},Rmat::Array{Float64,2},gamma::Float64,p::Array{Float64,1},memo::Dict{Pair{Vector{Int},Vector{Float64}},Float64},linememo::Dict{Pair{Vector{Int},Vector{Int}},Float64})
    if all(p.==1)
        return LK1(i,n,lambda,Rmat,gamma,memo,linememo)
    else
        ll=0.0
        p=fill(0.0,size(i)).+p
        for lat_i in VIterable(n,i)
            ll+=LK1(lat_i,n,lambda,Rmat,gamma,memo,linememo)+sum(logpdf.(Binomial.(lat_i,p),i))
            return ll
        end
    end
end

function LKnn(n::Array{Int,1},lambda::Array{Float64,1},Rmat::Matrix{Float64},nfactor::Array{Float64,1},gamma::Float64,memo::Dict{Pair{Vector{Int},Vector{Float64}},Float64})
	if haskey(memo,n=>nfactor) return memo[n=>nfactor] end
    if all(n.==0) return 0.0 end

	#foi=Rmat*k.*nfactor.+lambda

    deduct = [sum(log.(binomial.(n,k)))-sum((n-k).* (Rmat*k.*nfactor.+lambda) )+LKnn(k,lambda,Rmat,nfactor,gamma,memo) for k in VIterable(n) if k!=n]
    lpr=log1mexp(logsumexp(deduct))
	memo[n=>nfactor]=lpr
    return(lpr)
end

function LK1(i::Array{Int,1},n::Array{Int,1},lambda::Array{Float64,1},Rmat::Array{Float64,2},gamma::Float64,memo::Dict{Pair{Vector{Int},Vector{Float64}},Float64},linememo::Dict{Pair{Vector{Int},Vector{Int}},Float64})
    if haskey(linememo,n=>i) return linememo[n=>i] end
    Neff=(Rmat*n-diag(Rmat).*min.(n,1))/(sum(Rmat)/size(Rmat,1))
    nfactor=1.0./Neff.^gamma
	nfactor[isnan.(nfactor)].=0.0
    foi=Rmat*i.*nfactor.+lambda
    rtn=sum(log.(binomial.(n,i)))-sum((n-i).*foi)+LKnn(i,lambda,Rmat,nfactor,gamma,memo)
    linememo[n=>i]=rtn
    return rtn
end

function ll(i::Array{Int,2},n::Array{Int,2},lambda::Array{Float64,1},Rmat::Array{Float64,2},gamma:: Float64=0.0,p::Array{Float64,1}=[1.0])
	memo=Dict{Pair{Vector{Int},Vector{Float64}},Float64}() #memoize LKnn by n => nfactor
    linememo=Dict{Pair{Vector{Int},Vector{Int}},Float64}()
    if size(i)!=size(n) error("sizes of n and i don't match") end
    if !(size(i,1)==size(lambda,1)==size(Rmat,1)==size(Rmat,2)) error("the number of classes don't match") end
    llik=0.0
    for c in 1:size(n,2)
        llik+=LKlatent(i[:,c],n[:,c],lambda,Rmat,gamma,p,memo,linememo)
    end
    return llik
end
end


using ADCME
using PyPlot
using Random
using MAT
using DelimitedFiles
matplotlib.use("agg")

include("fwi_data.jl")
reset_default_graph()
tf.set_random_seed(123)
tf.random.set_random_seed(123)
θ0 = matread("results/weight$model.mat")["theta"][:]

μ = Variable(θ0)
ρ = Variable(zeros(length(θ0)))

if model=="step"
    σ = log(1 + exp(ρ))
else
    σ = 0.1*log(1 + exp(ρ))
end

# dt = 1e-10
dt = 1e-8
# dt = 0

function likelihood(z)
    θ = μ + z * σ

    R_ = constant(R)

    idx_var = pml+Int(round(0.1/Δx))

    inv_sigmoid(y) = @. - log(1.0/y - 1.0)

    # case = "fwi"
    case = "nnfwi"

    ## FC
    if case == "nnfwi"
    x0 = collect( range(0, 1, length=N-1-idx_var) )
    if model == "step"
        var= fc(x0, [30, 30, 30, 1], θ) * 0.1 + inv_sigmoid(1.3/2.0)
        var = -2.0*sigmoid(var)
    elseif model == "slop"
        var= fc(x0, [30, 30, 30, 1], θ) * 0.1
    end
    end

    cH = scatter_add(constant(cH_init), idx_var+1:N-1, var)

    cE = (cH[1:end-1]+cH[2:end])/2

    reg_TV = sum(abs(cH[2:end]-cH[1:end-1]))
    reg_smooth = sum((cH[1:end-2] + cH[3:end] - 2*cH[2:end-1])^2)

    function condition(i, E_arr, H_arr)
        i<=NT+1
    end

    function body(i, E_arr, H_arr)
        E = read(E_arr, i-1)
        H = read(H_arr, i-1)

        ΔH = cH * (E[2:end]-E[1:end-1])/Δx - σH*H
        H += ΔH * Δt
        ΔH = 1/(24Δx) * cH[2:N-2] * (-E[4:end] + 3E[3:end-1] - 3E[2:end-2] + E[1:end-3])
        H = scatter_add(H, 2:N-2, ΔH * Δt)

        ΔE = cE * (H[2:end]-H[1:end-1])/Δx - σE[2:end-1]*E[2:end-1] + R_[i] * Z
        E = scatter_add(E, 2:N-1, ΔE * Δt)
        ΔE = 1/(24Δx) * cE[2:end-1] * (-H[4:end] + 3H[3:end-1] - 3H[2:end-2] + H[1:end-3])
        E = scatter_add(E, 3:N-2, ΔE*Δt)

        i+1, write(E_arr, i, E), write(H_arr, i, H)
    end

    E_arr = TensorArray(NT+1)
    H_arr = TensorArray(NT+1)

    E_arr = write(E_arr, 1, zeros(N))
    H_arr = write(H_arr, 1, zeros(N-1))

    i = constant(2, dtype = Int32)

    _, E, H = while_loop(condition, body, [i, E_arr, H_arr])

    E = stack(E); E = set_shape(E, (NT+1, N))
    H = stack(H)

    loss = sum((E[:, idx_rcv] - E_true[:, idx_rcv])^2 * [1.0, 1e-3])

    lkh = dt*loss - sum((θ-μ)^2/σ^2) + sum((θ-θ0)^2)
    cH, lkh
end

z = tf.random_normal((8, length(θ0)), dtype = tf.float64)
# z = constant(randn(8, length(θ0)))
loss = constant(0.0)
for i = 1:size(z, 1)
    global cH, loss_ = likelihood(z[i])
    global loss += loss_
end
lr = placeholder(Float64, shape = [])
opt = AdamOptimizer(lr).minimize(loss)

sess = Session(); init(sess)

if dt>0.0
    global losses = []
    for i = 1:1000
        _, l = run(sess, [opt, loss], lr=>0.001)
        push!(losses, l)
        @info i, l
    end
    close("all")
    plot(losses)
    savefig("loss.png")
end
# error()
close("all")
cH, _ = likelihood(tf.random_normal((length(θ0),), dtype = tf.float64))
cH_est = zeros(length(cH), 100)
for i = 1:100
    cH_est[:,i] = run(sess, cH)
    plot(xH, cH_est[:,i], "--", color="gray", alpha = 0.5)
end
plot(xH, cH_true, "C0", label="True")
# plot(xH, cH0, "k--", label="Initial")
legend()

if dt≈0
    savefig("$(model)_prior.png")
else
    savefig("$model.png")
end


"""
    OILMM(fs, H)
An Orthogonal Instantaneous Linear Mixing Model (OILMM) -- a distribution over vector-
valued functions. Let `p` be the number of observed outputs, and `m` the number of latent
processes, then `H`, also known as the mixing matrix, is a `p x m` orthogonal matrix whose
column space spans the output space. The latent processes are represented an `Independent_MOGP`
as the latent processes remain decoupled.

# Arguments:
- f: a length-`m` vector of Gaussian process objects as an IndependentMOGP.
- H: a `p x m` orthogonal matrix representing a fixed basis of our p-dim target: h_1,...,h_m
"""
const OILMM = ILMM{<:IndependentMOGP, <:Orthogonal}

"""
    project(H, σ²)

Computes the projection `T` and `ΣT` given the mixing matrix and noise.
"""
function project(H::Orthogonal{Z}, σ²::Z) where {Z<:Real}
    U, S = H.U, H.S

    # Compute the projection of the data.
    T = sqrt(S) \ U'

    # Compute the projected noise, which is a matrix of size `size(Yproj)`.
    ΣT = diag(σ² * inv(S))

    return T, ΣT
end

"""
    rand(rng::AbstractRNG, fx::FiniteGP{<:OILMM})
Sample from the OILMM, including the observation noise.
Follows generative structure of model 2 from [1].
Follows the AbstractGPs.jl API.
See also `rand_latent`.
[1] - Bruinsma et al 2020.
"""
function AbstractGPs.rand(rng::AbstractRNG, fx::FiniteGP{<:OILMM})
    fs, H, σ², x = unpack(fx)

    # Obtain U and S
    U, S = H.U, H.S

    # Generate from the latent processes.
    X = reshape(reduce(vcat, map(f -> rand(rng, f(x)), fs.fs)), length(x), :)

    # Transform latents into observed space.
    F = vec((U * sqrt(S) * X')')

    # Generate iid noise and add to each output.
    return F .+ sqrt(noise_var(fx.Σy)) .* randn(rng, size(F))
end

# See AbstractGPs.jl API docs.
function AbstractGPs.mean_and_var(fx::FiniteGP{<:OILMM})
    fs, H, σ², x = unpack(fx)

    # Compute the marginals over the independent latents.
    fs_marginals = reduce(hcat, map(f -> AbstractGPs.marginals(f(x)), fs.fs))
    M_latent = mean.(fs_marginals)'
    V_latent = var.(fs_marginals)'

    # Obtain U and S
    U, S = H.U, H.S

    # Compute the means.
    M = U * sqrt(S) * M_latent

    # Compute the variances.
    V = abs2.(U * sqrt(S)) * (V_latent) .+ σ²

    # Package everything into independent Normal distributions.
    return collect(vec(M')), collect(vec(V'))
end

# See AbstractGPs.jl API docs.
function AbstractGPs.logpdf(fx::FiniteGP{<:OILMM}, y::AbstractVector{<:Real})
    fs, H, σ², x = unpack(fx)

    # Projection step.
    Y = reshape_y(y, length(x))
    T, ΣT = project(H, σ²)
    Ty = T*Y

    # Latent process log marginal likelihood calculation.
    y_rows = collect(eachrow(Ty))
    ΣT_rows = collect(eachrow(repeat(ΣT, 1, size(Ty, 2))))
    lmls_latents = map((f, s, y) -> logpdf(f(x, collect(s)), collect(y)), fs.fs, ΣT_rows, y_rows)

    return sum(lmls_latents) + regulariser(H, σ², Y)
end

"""
    regulariser(fx, y)

Computes the regularisation term of the logpdf.
See e.g. appendix A.4 of [1] - Bruinsma et al 2020.
"""
function regulariser(
    H::Orthogonal{T},
    σ²::T,
    Y::AbstractMatrix{T},
) where {T<:Real}
    U, S = H.U, H.S

    n = size(Y, 2)
    p, m = size(U)

    return -(n * (logdet(S) + (p - m) * log(2π * σ²)) +
        sum(abs2, (I - U * U') * Y) / σ²) / 2
end

# See AbstractGPs.jl API docs.
function AbstractGPs.posterior(fx::FiniteGP{<:OILMM}, y::AbstractVector{<:Real})
    fs, H, σ², x = unpack(fx)

    # Projection step.
    Y = reshape_y(y, length(x))
    T, ΣT = project(H, σ²)
    Ty = T*Y

    # Condition each latent process on the projected observations.
    y_rows = collect(eachrow(Ty))
    ΣT_rows = collect(eachrow(repeat(ΣT, 1, size(Ty, 2))))
    fs_posterior = map(
        (f, s, y) -> AbstractGPs.posterior(f(x, collect(s)), collect(y)),
        fs.fs,
        ΣT_rows,
        y_rows,
    )
    return ILMM(independent_mogp(fs_posterior), H)
end


using LinearAlgebra

# Define a generic linear operator with preallocation
function preallocated_LinearOperator(A)
  (n, m) = size(A)
  Ap = zeros(n)
  Atq = zeros(m)
  return LinearOperator(n, m, false, false, p -> mul!(Ap, A, p),
                        q -> mul!(Atq, transpose(A), q), q -> mul!(Atq, transpose(A), q))
end

# Variants where A is a matrix without specific properties
for fn in (:cgls, :cgne, :craig, :craigmr, :crls, :crmr, :lslq, :lsmr, :lsqr, :dqgmres, :diom, :cgs)
  @eval begin
    $fn(A :: AbstractMatrix{TA}, b :: AbstractVector{Tb}, args...; kwargs...) where {TA, Tb <: Number} =
      $fn(preallocated_LinearOperator(A), convert(Vector{Float64}, b), args...; kwargs...)
  end
end

# Define a symmetric linear operator with preallocation
function preallocated_symmetric_LinearOperator(A)
  (n, m) = size(A)
  Ap = zeros(n)
  return LinearOperator(n, m, true, true, p -> mul!(Ap, A, p))
end

# Variants where A must be symmetric
for fn in (:cg_lanczos, :cg_lanczos_shift_seq, :cg, :cr, :minres, :symmlq)
  @eval begin
    $fn(A :: AbstractMatrix{TA}, b :: AbstractVector{Tb}, args...; kwargs...) where {TA, Tb <: Number} =
      $fn(preallocated_symmetric_LinearOperator(A), convert(Vector{Float64}, b), args...; kwargs...)
  end
end


import DoubleFloats: Double64

issquare(A::AbstractMatrix) = isequal(size(A)...)

"""
    row_scaling!(A, D::AbstractVector, tol=maximum(size(A))^2 * eps(real(eltype(A))))

Divide each row of `A` by its ∞-norm. This minimizes the ∞ condition number
for all possible scalings. In order to avoid scaling near zero rows back to
1 we only scale a row if its ∞-norm is larger than `tol`.
"""
function row_scaling!(A::AbstractMatrix{T}, D::AbstractVector, tol=max(size(A)...)^2 * eps(real(T))) where {T}
    D .= 0.0
    d_max = convert(real(T), -Inf)
    @inbounds for j=1:size(A, 2), i=1:size(A, 1)
        D[i] = @fastmath max(D[i], abs(A[i,j]))
        d_max = @fastmath max(d_max, D[i])
    end

    # In the case that the row norms are all less than 1 we have a lower bound
    # the maximal row norm
    # To avoid that we scale a zero matrix, we need a tol for wich we don't scale any more
    bound = clamp(d_max, tol, 1.0)
    D .= max.(D, bound)

    @inbounds for j=1:size(A, 2), i=1:size(A, 1)
        A[i, j] /= D[i]
    end
    A
end

@enum ActiveFactorization begin
    QR_FACTORIZATION
    LU_FACTORIZATION
end

"""
    Jacobian{T, <:Factorization}

Data structure holding a Jacobian matrix `J` with it's factorization.
Additional contains a vector `D` of scalings which row equilibarate `J`.

    Jacobian(J)

Setup a Jacobian with Jacobian matrices of the form `J`.
"""
mutable struct Jacobian{T}
    J::Matrix{T} # Jacobian matrix
    D::Vector{Float64} # Inverse of scaling factors D
    # lu is only there for square systems
    lu::Union{Nothing, LinearAlgebra.LU{T,Matrix{T}}} # LU Factorization of D * J
    qr::LinearAlgebra.QRPivoted{T,Matrix{T}} # Pivoted QR Factorization of D * J
    active_factorization::ActiveFactorization
    b::Vector{T} # Vector for overdetermined systems
    r::Vector{T} # Vector for iterative refinement
    perm::Vector{Int} # work copy of the permutation in the qr
    qr_work::Vector{ComplexF64} # preallocated memory for pivoted qr
    qr_rwork::Vector{Float64} # preallocated memory for pivoted qr
    ormqr_work::Vector{ComplexF64}
    # Numerical Informations
    corank::Int # The numerical corank
    corank_proposal::Int
    cond::Float64 # Estimate of the Jacobian
    # The relative number of digits lost during the solution of the linear systems
    # in Newton's method. See `solve_with_digits_lost!` in utilities/linear_algebra.jl
    # for how this is computed.
    digits_lost::Union{Nothing,Float64}
end

function Jacobian(A::AbstractMatrix, corank::Int=0)
    m, n = size(A)
    lu = m ≠ n ? nothing : LinearAlgebra.lu!(Random.randn!(complex.(float.(A))))

    work = Vector{ComplexF64}(undef, 1)
    ormqr_work = Vector{ComplexF64}(undef, 1)
    rwork = Vector{Float64}(undef, 2n)
    jpvt = zeros(BlasInt, n)
    B = ComplexF64.(A)
    tau = similar(B, min(m, n))
    geqp3!(B, jpvt, tau, work, rwork)
    qr = LinearAlgebra.QRPivoted(B, tau, jpvt)
    active_factorization = lu === nothing ? QR_FACTORIZATION : LU_FACTORIZATION
    J = copy(B)
    D = ones(m)
    r = similar(J, size(J, 1))
    b = copy(r)
    perm = zeros(Int, length(qr.p))
    cond = 1.0
    digits_lost = nothing
    Jacobian(J, D, lu, qr, active_factorization, b, r, perm, work, rwork, ormqr_work,
        corank, 0, cond, digits_lost)
end

function update_rank!(Jac::Jacobian)
    Jac.corank = Jac.corank_proposal
    Jac
end

function reset!(jacobian::Jacobian)
    jacobian.cond = 1.0
    jacobian.corank = 0
    jacobian.corank_proposal = 0
    jacobian.digits_lost = nothing
    jacobian
end

"""
    update!(jacobian, J, corank::Int=jacobian.corank)

Update the `jacobian` struct with the new Matrix `J`.
"""
function update!(Jac::Jacobian, J::AbstractMatrix; update_infos::Bool=false)
    copyto!(Jac.J, J)
    updated_jacobian!(Jac; update_infos=update_infos)
    Jac
end

"""
    update_jacobian!(jacobian)

This computes a factorization of the stored Jacobian matrix.
Call this instead of `update!` if `jacobian.J` got updated.
"""
function updated_jacobian!(Jac::Jacobian{ComplexF64}; update_infos::Bool=false) where {T}
    if !update_infos && issquare(Jac.J) && Jac.corank == 0
        Jac.lu !== nothing || return Jac
        copyto!(Jac.lu.factors, Jac.J)
        lu!(Jac.lu.factors, nothing, Jac.lu.ipiv)
        Jac.active_factorization = LU_FACTORIZATION
    else
        copyto!(Jac.qr.factors, Jac.J)

        m, n = size(Jac.J)
        if m == n
            # only apply row scaling for square matrices since
            # otherwise we change the problem
            row_scaling!(Jac.qr.factors, Jac.D, 1)
        end
        # this computes a pivoted qr factorization, i.e.,
        # qr!(Jac.qr.factors, Val(true)) but without allocating new memory
        geqp3!(Jac.qr.factors, Jac.qr.jpvt, Jac.qr.τ, Jac.qr_work, Jac.qr_rwork)

        ε = max(n,m) * eps() * 10
        # check rank 0
        rnm = min(n,m)
        r₁ = abs(real(Jac.qr.factors[1,1]))
        corank = 0
        if r₁ < ε
            corank = rnm
            Jac.cond = inv(ε)
        else
            for i in 2:rnm
                rᵢ = abs(real(Jac.qr.factors[i,i]))
                if ε * r₁ > rᵢ
                    corank = rnm - i + 1
                    break
                end
            end
            Jac.cond = max(1.0, r₁) / abs(real(Jac.qr.factors[rnm, rnm]))
        end
        Jac.corank_proposal = corank

        Jac.active_factorization = QR_FACTORIZATION
    end
    Jac
end

"""
    solve!([x, ], Jac::Jacobian, b)

Solve `Jac.J * x = b` inplace. Assumes `Jac` contains already the correct factorization.
This stores the result in `x`.
"""
function solve!(x, Jac::Jacobian, b::AbstractVector; update_digits_lost::Bool=false, refinement_step::Bool=false)
    if Jac.active_factorization == LU_FACTORIZATION
        lu = Jac.lu
        if lu !== nothing
            lu_ldiv!(x, lu, b)
        end
    elseif Jac.active_factorization == QR_FACTORIZATION
        if issquare(Jac.J)
            # we applied row scaling for square matrices
            @inbounds for i in eachindex(Jac.b)
                Jac.b[i] = b[i] / Jac.D[i]
            end
            custom_ldiv!(x, Jac.qr, Jac.b, Jac.corank, Jac.perm, Jac.ormqr_work )
        else
            Jac.b .= b
            custom_ldiv!(x, Jac.qr, Jac.b, Jac.corank, Jac.perm, Jac.ormqr_work )
        end
    end
    if issquare(Jac.J) && Jac.corank == 0 && update_digits_lost && !refinement_step
        norm_x = euclidean_norm(x)
        norm_δx = iterative_refinement_step!(x, Jac, b)
        Jac.digits_lost = log₁₀(norm_δx / eps(norm_x))
    elseif !issquare(Jac.J)
        Jac.digits_lost = nothing
    end
    x
end
solve!(A::Jacobian, b) = solve!(b, A::Jacobian, b)

const LA = LinearAlgebra

# This is an adoption of the Julia implementation of ldiv! for QRPivoted
function custom_ldiv!(x, A::LinearAlgebra.QRPivoted{ComplexF64},
                      b::Vector{ComplexF64}, corank::Int,
                      perm::Vector{Int}, ormqr_work::Vector{ComplexF64})
    mA, nA = size(A.factors)
    nr = min(mA,nA)
    nrhs = length(b)
    if nr == 0
        return x
    end
    if corank == nr
        fill!(x, zero(eltype(x)))
        return x
    end
    rank = nr - corank

    if rank == nA == nr
        # The following is equivalent to LA.lmul!(LA.adjoint(A.Q), b) but
        # we can also pass the preallocated work vector
        ormqr!('L', 'C', A.factors, A.τ, b, ormqr_work)
        ldiv_upper!(A.factors, b; singular_exception=false)
        @inbounds for i in 1:nr
            x[i] = b[i]
            perm[i] = A.p[i]
        end
        Base.invpermute!!(x, perm)
    else
        C, τ = LA.LAPACK.tzrzf!(A.factors[1:rank,:])
        LA.ldiv!(LA.UpperTriangular(C[1:rank,1:rank]),view(LA.lmul!(LA.adjoint(A.Q), b), 1:rank))
        for i in 1:rank
            x[i] = b[i]
        end
        for i in rank+1:nA
            x[i] = zero(eltype(x))
        end
        LA.LAPACK.ormrz!('L', eltype(x)<:Complex ? 'C' : 'T', C, τ, reshape(x, length(x), 1))
        invpermute!(x, A.p)
    end
    return x
end

"""
    iterative_refinement_step!(x, A, b, fac, T, r)

Apply one step of iterative refinement where the residual is computed with precision `T`.
Returns the euclidean norm of the update.
"""
function iterative_refinement_step!(x, Jac::Jacobian, b, ::Type{T}=eltype(x)) where T
    residual!(Jac.r, Jac.J, x, b, T)
    δx = solve!(Jac.r, Jac, Jac.r; refinement_step=true)
    norm_δx = euclidean_norm(δx)
    @inbounds for i in eachindex(x)
        x[i] = convert(T, x[i]) - convert(T, δx[i])
    end

    return norm_δx
end


using SimpleBirthDeathProcess;
using DelimitedFiles;
using Plots;

N = 1_000;
F = Float64;
I = 25;
J = 35;
T = 2;
L = F(1);
M = F(3);

# values obtained with Maple 2018.2
correct_values = readdlm("log_trans_prob_values_univariate.csv", ',', BigFloat);

μset = [M * F(n) / N for n = 1:N];
logp = [trans_prob(I, J, T, [L; F(μ)]) for μ = μset];

relative_error = abs.(1 .- logp ./ correct_values);

pl = plot(μset, relative_error, seriestype=:line, colour=:black,
          yscale=:log10, xlab="\\mu", ylab="Relative error",
          legend=false, size=(1920, 1080));

savefig(pl, "Figure_3.svg")


"""
    function distribution()
---
This function is to test how to simulate correlated traits.
"""
function distribution(;nqtl = 1000)
    @info "Simulate a larege population"
    # to simulate
    #base = sim_base(20000, 24, 1e8, 4000, maf=0.01)
    #serialize("dat/run/large-base.ser", base)
    # to use simulate results above
    base = deserialize("dat/run/large-base.ser")
    ds = [Gamma(0.4),
          Laplace(),
          Normal(),
          Uniform(),
          Weibull(),
          Chi(1)]
    label = [L"\Gamma(0.4)", "Laplace()", L"N(0, 1)", L"U(0, 1)", "Weibull()", L"\chi^2(1)"]
    for i in 1:length(ds)
        qtl = sim_QTL(base, nqtl, d=ds[i])[1]
        qgt = qtl_gt(base[:hap], qtl.pos)
        tbv = qgt'qtl.effect
        histogram(tbv, normalize=true, label=label[i], dpi=300)
        savefig("dat/run/$i.pdf")
    end
end


using CompScienceMeshes, BEAST
o, x, y, z = euclidianbasis(3)

sol = 5.0
#Δt, Nt = 0.08/sol,400
Δt, Nt = 0.120000123/sol,400

#D, Δx = 1.0, 0.25
D, Δx = 1.0, 0.1
#Γ = meshsphere(D, Δx)
Γ = meshrectangle(D,D,Δx,3)
#γ = boundary(Γ)
γ1 = meshsegment(D,D,3)
γ2 = CompScienceMeshes.translate(γ1, point(0,D,0))

#X = raviartthomas(Γ,weld(γ1,γ2))
X = rt_ports(Γ,(γ1,γ2))
X.fns = X.fns[1:281]
numfunctions(X)

T = timebasisshiftedlagrange(Δt, Nt, 3)
U = timebasisdelta(Δt, Nt)

V = X ⊗ T
W = X ⊗ U

duration, delay, amplitude = 8.0/sol, 12.0/sol, 1.0
gaussian = creategaussian(duration, delay, duration)

direction, polarisation = z, x
E = planewave(polarisation, direction, derive(gaussian), sol)
T = MWSingleLayerTDIO(sol,-1/sol,-sol,2,0)

#Z = assemble(T,W,V);

@hilbertspace j
@hilbertspace j′
tdefie = @discretise T[j′,j] == -1E[j′]   j∈V  j′∈W
xefie = solve(tdefie)

using Plotly
# include(Pkg.dir("CompScienceMeshes","examples","matlab_patches.jl"))

Xefie, Δω, ω0 = fouriertransform(xefie, Δt, 0.0, 2)
ω = collect(ω0 + (0:Nt-1)*Δω)
_, i1 = findmin(abs.(ω.-sol))

ω1 = ω[i1]
ue = Xefie[:,i1] / fouriertransform(gaussian)(ω1)

using LinearAlgebra
fcre, geo = facecurrents(ue, X)
t2 = patch(geo, real.(norm.(fcre)))
Plotly.plot(t2)

fcr_td, geo = facecurrents(xefie[:,end], X)
# fcr_ch, geo = facecurrents(xefie[:,end], divergence(X))

t3 = patch(Γ,real.(norm.(fcr_td)))
Plotly.plot(t3)
# jmatlab_quiver(Γ, fcr_td)


module TestTikZWiringDiagrams

using Test

using Catlab.Doctrines
using Catlab.WiringDiagrams
using Catlab.Graphics
import Catlab.Graphics: TikZ
import Catlab.Graphics.TikZWiringDiagrams.Defaults

# We can't test that the pictures look right, but we can test that they exist!
is_pic(obj) = isa(obj, TikZ.Picture)

A, B = Ob(FreeCategory, :A, :B)
f, g = Hom(:f, A, B), Hom(:g, B, A)
@test is_pic(to_tikz(f))
@test is_pic(to_tikz(compose(f,g)))
@test is_pic(to_tikz(compose(id(A), f, id(B))))

A, B = Ob(FreeDaggerCategory, :A, :B)
f, g = Hom(:f, A, B), Hom(:g, B, A)
@test is_pic(to_tikz(dagger(f)))
@test is_pic(to_tikz(dagger(compose(f,g))))

A, B = Ob(FreeSymmetricMonoidalCategory, :A, :B)
f, g = Hom(:f, A, B), Hom(:g, B, A)
@test is_pic(to_tikz(otimes(f,g)))
@test is_pic(to_tikz(compose(otimes(f,g),braid(B,A))))

A, B = Ob(FreeBiproductCategory, :A, :B)
f, g = Hom(:f, A, B), Hom(:g, B, A)
@test is_pic(to_tikz(compose(mcopy(A), otimes(f,f), mmerge(B))))
@test is_pic(to_tikz(compose(create(A), f, g, delete(A))))

A, B = Ob(FreeCompactClosedCategory, :A, :B)
f, g = Hom(:f, A, B), Hom(:g, B, A)
h = Hom(:h, otimes(A,B), otimes(A,B))
@test is_pic(to_tikz(dunit(A)))
@test is_pic(to_tikz(dcounit(A)))
trace = compose(
  f,
  otimes(dunit(A), id(B)),
  otimes(id(dual(A)), h),
  otimes(braid(dual(A),A), id(B)),
  otimes(dcounit(A), id(B)),
  g,
)
@test is_pic(to_tikz(trace))

end