"""
    prediction( sys::StateSpace, N::Integer )

Form the prediction matrix ``Γ`` that computes the states of the linear system `sys` over the horizon of length `N`
by using all the inputs across the entire horizon.

# Extended help

This matrix consists of the blocks ``A^i B`` where ``i`` is the diagonal number, forming the matrix
```math
⎡ B   •   •  •  ⋯⎤
  ⎢ AB  B   •  •  ⋯⎥
  ⎢ A²B AB  B  •  ⋯⎥
  ⎢ A³B A²B AB B  ⋯⎥
  ⎣ ⋮   ⋮   ⋮  ⋮ ⋱ ⎦
```

It is used to compute the component of the state vector for a horizon that is caused by the system inputs
during the horizon using the equation ``x = Γu``.

The full state vector for a horizon can be computed as ``x = Γu + Φx₀``, where ``Φ`` is the initial condition
propagation matrix from [`initialpropagation`](@ref).
"""
function prediction( sys::StateSpace, N::Integer )
    if( N < 1 )
        throw( DomainError( N, "Horizon length must be greater than 1" ) )
    end

    nu = ninputs( sys )
    nx = nstates( sys )
    mt = eltype( sys.A )

    mainDiag  = kron( I(N), sys.B )

    otherDiag = kron( I(N), to_matrix( mt, I, nx ) )

    D = zeros( mt, N, N )
    D[diagind( D, -1 )] .= 1

    otherDiag = otherDiag + kron( D, -sys.A )

    Γ = otherDiag\mainDiag

    return BlockArray{ mt }( Γ, [nx for i = 1:N], [nu for i = 1:N] )
end

"""
    initialpropagation( sys::StateSpace, N::Integer; ive::Bool = false )

Form the matrix ``Φ`` that propagates the initial condition ``x₀`` in the linear system `sys` across the horizon `N`.
When `ive` is true, the initial condition ``x₀`` is embedded in the resulting state vector

# Extended help

When `ive` is false, this matrix consists of a column of the blocks ``A^i`` where ``i`` is the row number, forming the matrix
```math
⎡ A  ⎤
  ⎢ A² ⎥
  ⎢ A³ ⎥
  ⎢ A⁴ ⎥
  ⎣ ⋮  ⎦
```

When `ive` is true, this matrix consists of a column of the blocks ``A^i-1`` where ``i`` is the row number, forming the matrix
```math
⎡ I  ⎤
  ⎢ A² ⎥
  ⎢ A² ⎥
  ⎢ A³ ⎥
  ⎣ ⋮  ⎦
```

It is used to compute the component of the state vector for a horizon that is caused by the intial state of the system using
the equation ``x = Φx₀``.

The full state vector for a horizon can be computed as ``x = Γu + Φx₀``, where ``Γ`` is the prediction matrix
from [`prediction`](@ref).
"""
function initialpropagation( sys::StateSpace, N::Integer; ive::Bool = false )
    if( N < 1 )
        throw( DomainError( N, "Horizon length must be greater than 1" ) )
    end

    sub = ive ? 1 : 0

    blockGen = (sys.A^(i-sub) for i=1:N)

    # The mortar function requires the passed array to have ndims=2, so reshape to
    # have an array with 2 dimensions even though there is only one column
    blocks = reshape( collect( blockGen ),  (N, 1) )

    return mortar( blocks )
end

"""
    hessian( clqr::ConstrainedTimeInvariantLQR, ::Type{PrimalFullCondensing} )

Form the complete Hessian for the constrained time-invariant LQR problem `clqr` using the fully
condensed form.
"""
function hessian( clqr::ConstrainedTimeInvariantLQR, ::Type{PrimalFullCondensing} )
    sysₖ = getsystem( clqr, prestabilized = true )
    N    = clqr.N
    nₓ   = nstates( sysₖ )
    nᵤ   = ninputs( sysₖ )

    Iₙ = 1.0*I(N)
    Γₖ = prediction( sysₖ, N )

    Q = clqr.Q
    S = clqr.S
    K = clqr.K
    R = clqr.R

    Qₖ = Q - K'*S' - S*K + K'*R*K
    Sₖ = S - K'*R

    Q̄ₖ = blockkron( Iₙ, Qₖ )
    S̄ₖ = blockkron( Iₙ, Sₖ )
    R̄  = blockkron( Iₙ, R  )

    Q̄ₖ[Block( N, N )] = clqr.P
    S̄ₖ[Block( N, N )] = zeros( eltype( clqr.S ), size( clqr.S ) ) - K'*R

    # Pass the result through the Hermitian type to cleanup numerical errors that can occur
    # making it non-symmetric (they are on the order of 10e-17)
    return Hermitian( Γₖ'*Q̄ₖ*Γₖ + Γₖ'*S̄ₖ + S̄ₖ'*Γₖ + R̄  )
end

function hessian_shifted( clqr::ConstrainedTimeInvariantLQR, ::Type{PrimalFullCondensing} )
    sysₖ = getsystem( clqr, prestabilized = true )
    mt   = eltype( sysₖ.A )
    N    = clqr.N
    N₁   = N+1
    nₓ   = nstates( sysₖ )
    nᵤ   = ninputs( sysₖ )

    Iₙ  = 1.0*I(N)
    Iₙ₁ = 1.0*I(N₁)

    # Form the prediction matrix with the initial value embedded
    Γₖ = prediction( sysₖ, N₁ )
    Γ̃ₖ = BlockArray{ mt }( zeros( mt, nₓ*N₁, nᵤ*N ), [nₓ for i = 1:N₁], [nᵤ for i = 1:N] )
    [Γ̃ₖ[Block(i, j)] = Γₖ[Block(i-1, j)] for i=2:N₁, j=1:N]

    Q̄ₖ = blockkron( Iₙ₁,  clqr.Q )
    S̄  = blockkron(  Iₙ,  clqr.S )
    R̄  = blockkron(  Iₙ,  clqr.R )
    K̄  = blockkron(  Iₙ, -clqr.K )

    Q̄ₖ[Block( N₁, N₁ )] = clqr.P
    S̄ = vcat( S̄, zeros( eltype( clqr.S ), (nₓ, N*nᵤ ) ) )
    K̄ = hcat( K̄, zeros( eltype( clqr.K ), (N*nᵤ, nₓ ) ) )

    # Pass the result through the Hermitian type to cleanup numerical errors that can occur
    # making it non-symmetric (they are on the order of 10e-17)
    return Hermitian( Γ̃ₖ'*(Q̄ₖ + K̄'*R̄*K̄ + S̄ *K̄ + K̄'*S̄' )*Γ̃ₖ + Γ̃ₖ'*(K̄' *R̄ + S̄ ) + (R̄ *K̄ + S̄')*Γ̃ₖ + R̄  )
end

"""
    linearcoefficients( clqr::ConstrainedTimeInvariantLQR, ::Type{PrimalFullCondensing} )

Form the vector of coefficients for the linear term of the fully condensed quadratic QP for the
CLQR problem defined by `clqr`.
"""
function linearcoefficients( clqr::ConstrainedTimeInvariantLQR, ::Type{PrimalFullCondensing} )
    sysₖ = getsystem( clqr, prestabilized = true )
    N    = clqr.N
    nx   = nstates( sysₖ )
    nu   = ninputs( sysₖ )

    Iₙ = 1.0*I(N)
    Γₖ = prediction( sysₖ, N )
    Φₖ = initialpropagation( sysₖ, N )

    Q̄ₖ = blockkron( Iₙ,  clqr.Qₖ )
    S̄  = blockkron( Iₙ,  clqr.S  )
    R̄  = blockkron( Iₙ,  clqr.R  )
    K̄  = blockkron( Iₙ, -clqr.K  )

    Q̄ₖ[Block( N, N )] = clqr.P
    S̄[Block( N, N )]  = zeros( eltype( clqr.S ), size( clqr.S ) )

    # Temporary product used in two places
    T = 1.0*I(N*nu) + Γₖ'*K̄'

    return ((Γₖ'*S̄ + T*R̄')*K̄ + T*S̄' + Γₖ'*Q̄ₖ')*Φₖ
end

function inequalityconstraints( clqr::ConstrainedTimeInvariantLQR, ::Type{PrimalFullCondensing} )
    sysₖ = getsystem( clqr, prestabilized = true )

    mt   = eltype( sysₖ.A )
    N    = clqr.N
    nx   = nstates( sysₖ )
    nu   = ninputs( sysₖ )
    Iₙ   = 1.0*I(N)

    # Get the condensed initial propagation matrix
    Φ̃ₖ = initialpropagation( sysₖ, N, ive = true )

    # Form the prediction matrix with the initial value embedded
    Γₖ = prediction( sysₖ, N )
    Γ̃ₖ = BlockArray{ mt }( zeros( mt, nx*N, nu*N ), [nx for i = 1:N], [nu for i = 1:N] )
    [Γ̃ₖ[Block(i, j)] = Γₖ[Block(i-1, j)] for i=2:N, j=1:N]

    # Create the component matrices
    Ē  = blockkron( Iₙ, clqr.E )
    F̄  = blockkron( Iₙ, clqr.F )
    K̄  = blockkron( Iₙ, clqr.K  )

    # Common matrix
    # The minus sign here comes from the fact we have a negative feedback controller (u = -Kx)
    T = Ē - F̄ *K̄

    # Compute the coefficient matrix of the linear inequality constraints
    G = T*Γ̃ₖ + F̄

    # Compute the initial state matrix for the RHS of the constraints
    L = -T*Φ̃ₖ;

    # Compute the constant vector for the RHS
    g = blockkron( ones( Float64, (N) ), clqr.g )

    return G, L, g
end

function equalityconstraints( ::ConstrainedTimeInvariantLQR, ::Type{PrimalFullCondensing} )
    # The fully condensed form does not have any equality constraints from the problem
    return nothing
end


using Test
using Quaternions

test_associative(x, y, z, ⊗) = @test (x ⊗ y) ⊗ z ≈ x ⊗ (y ⊗ z)
test_commutative(x, y, ⊗) = @test x ⊗ y ≈ y ⊗ x
test_inverse(x, eins, ⊗, inv) = (@test x ⊗ inv(x) ≈ eins; @test inv(x) ⊗ x ≈ eins)
test_neutral(x, eins, ⊗) = (@test x ⊗ eins ≈ x; @test eins ⊗ x ≈ x)
test_monoid(x, y, z, ⊗, eins) = (test_associative(x, y, z, ⊗); test_neutral(x, eins, ⊗))
test_group(x, y, z, ⊗, eins, inv) = (test_monoid(x, y, z, ⊗, eins);test_inverse(x, eins, ⊗, inv))
test_multiplicative(x, y, ⊗, f) = @test f(x ⊗ y) ≈ f(x) ⊗ f(y)

include("test_Quaternion.jl")


using Test

@testset "(32)-1->(22) fixed-direction signaling regression" begin

using BellScenario: Degeneracy, LocalPolytope, ConvexPolytope

# """
# maxwell_chitambar_inequalities()
#     Uses input/output relabelings to construct the complete set of bell inequalities
#     from a canonical subset of base inequalities.
#
# Outputs:
#     ineqs: Array, contains inequality vectors.
# """
function maxwell_chitambar_inequalities()

    joint_pos_ineq = [0 0 0 0 -1 0 0 0 0 0 0 0 0 0 0 0]
    marginal_pos_ineq = [0 -1 0 0 1 0 0 0 0 0 0 0 0 0 0 0]
    I_ineq = [-1 -1 0 0 0 0 -1 1 1 0 -1 -1 0 0 1 0]
    II_ineq = [-1 0 0 0 -1 -1 -1 1 1 0 -1 -1 0 0 1 0]
    III_ineq = [-1 -1 -1 0 0 0 1 1 -1 1 -1 -1 0 0 1 -1]
    IV_ineq = [-1 0 0 0 -1 -1 -1 1 1 0 -1 -1 -1 1 1 0]
    V_ineq = [-1 -1 0 0 1 1 -1 -1 -1 1 0 0 -1 -1 1 -1]
    VI_ineq = [-1 -3 0 0 2 2 -2 1 0 -1 -1 -1 -1 0 1 -2]
    VII_ineq = [-1 -3 0 0 2 2 -2 1 0 -1 -1 -1 -2 1 1 -2]
    VIII_ineq = [-1 -3 0 0 2 2 -2 1 1 -2 -1 -1 -2 1 1 -2]

    ineqs = []

    # TODO: construct relabelngs in a way that is inclusive, but does not overshoot
    # by a factor of 4.5x duplicates.
    in_relabels = Degeneracy.bipartite_input_relabels((3,2),(2,2),"fixed-direction")
    (α_out_relabels, β_out_relabels) = Degeneracy.bipartite_output_relabels((3,2),(2,2),"fixed-direction")
    for in_relabel in in_relabels
        for α_out_relabel in α_out_relabels
            for β_out_relabel in β_out_relabels

                relabel = in_relabel*α_out_relabel*β_out_relabel

                push!(ineqs, joint_pos_ineq*relabel)
                push!(ineqs, marginal_pos_ineq*relabel)
                push!(ineqs, I_ineq*relabel)
                push!(ineqs, II_ineq*relabel)
                push!(ineqs, III_ineq*relabel)
                push!(ineqs, IV_ineq*relabel)
                push!(ineqs, V_ineq*relabel)
                push!(ineqs, VI_ineq*relabel)
                push!(ineqs, VII_ineq*relabel)
                push!(ineqs, VIII_ineq*relabel)
            end
        end
    end

    # TODO: fix so that we do not need to de-dupe
    unique(ineqs)
end

@testset "verifying polytope" begin

    dir = "./test/regression/files/"

    vertices = LocalPolytope.vertices((3,2),(2,2), bits=1)

    @test size(vertices) == (320,)
    @test size(vertices[1]) == (15,1)

    constraints = ConvexPolytope.facet_constraints(vertices, dir=dir)
    @test constraints[2] == []
    @test constraints[3] == []

    ineq_matches = maxwell_chitambar_inequalities()

    inequalities = constraints[1]
    @test issetequal(inequalities, ineq_matches)
end

end


using StanModels

df = filter(row -> row[:age] >= 18,
  CSV.read(joinpath(@__DIR__, "..", "..", "data", "Howell1.csv"), delim=';'))

max_height = maximum(df[!, :height])
min_height = minimum(df[!, :height])

m4_2s = "
data {
  int N;
  real h[N];
}
parameters {
  real<lower=0> sigma;
  real<lower=$(min_height),upper=$(max_height)> mu;
}
model {
  // Priors for mu and sigma
  mu ~ normal(178, 0.1);
  sigma ~ uniform( 0 , 50 );

  // Observed heights
  h ~ normal(mu, sigma);
}
";

sm = SampleModel("m4.2s", m4_2s);

m4_2_data = Dict("N" => size(df, 1), "h" => df[!, :height]);

(sample_file, log_file) = stan_sample(sm, data=m4_2_data);

if !(sample_file == nothing)
  chn = read_samples(sm)
  chn = set_names(chn, Dict("mu" => "μ", "sigma" => "σ"))
  describe(chn)
end


@with_kw mutable struct BatchRegularizer
    buffers # Datasets to use in the regularization
    𝒟s = nothing# Batches of those datasets
    batch_size::Int # Batch size of the data sets
    λ::Float32 = 1f0# regularization coefficients
    loss # Loss function that takes args (π, 𝒟)
end

value_regularization(π, 𝒟) = Flux.mse(value(π, 𝒟[:s]), 𝒟[:value])
action_regularization(π, 𝒟) = Flux.mse(action(π, 𝒟[:s]), 𝒟[:a])
action_value_regularization(π, 𝒟) = begin
    v = value(π, 𝒟[:s], 𝒟[:a])
    v isa Tuple && (v = v[1])
    Flux.mse(v, 𝒟[:value])
end

function (R::BatchRegularizer)(π)
    # sample a random batch for each buffer
    ret0 = false
    ignore() do
        if any([length(b) == 0 for b in R.buffers])
            ret0 = true
        else
            isnothing(R.𝒟s) && (R.𝒟s = [buffer_like(b, capacity=R.batch_size, device=device(π)) for b in R.buffers])
            for (𝒟, buffer) in zip(R.𝒟s, R.buffers)
                rand!(𝒟, buffer)
            end
        end
    end

    # Return the mean
    ret0 ? 0f0 : R.λ*mean([R.loss(π, 𝒟) for 𝒟 in R.𝒟s])
end


# Generate random channels -

# Copyright Robin Harper 2016-2020

using Distributions
using Random
using LinearAlgebra
using SparseArrays

# Note this is for reproducibility --- for multiple runs or spreadsheets you will need some
# other way to set.
#srand()
Random.seed!

_d=Normal()

#    Based on 1707.06926 (Rudnicki et al) and some others - see workbook for more complete description

"""
    returns the fidelity of the channel (in PauliLiouville basis)
"""
function fidelity(channel)
    diag = tr(channel)-1
    d = sqrt(size(channel)[1])
    p = diag/(d^2-1)
    return ((d-1)*p+1)/d
    #(1+1/3*(channel[2,2]+channel[3,3]+channel[4,4]))/2
end

"""
    returns the unitarity of the channel (PauliLiouville basis)
"""
function unitarity(noise)
    d = sqrt(size(noise)[1])
    return (1/(d^2-1))*tr(noise[2:end,2:end]'*noise[2:end,2:end])
end

"""
    A given fidelity has a minimum unitarity that it can have
    (so perfect fidelity = 1 = unitarity)
    This shows as a percentage how much the unitarity of the channel
    lies between the minimum and the maximum (always 1)
"""
function unitarityPercent(ch)
    fid = fidelity(ch)
    uni = unitarity(ch)
    d = sqrt(size(ch)[1])
    r = 1-fid
    minU = (1-(d*r)/(d-1))^2
    return (uni-minU)/(1-minU)
end

"""
    Simple function that takes a normal distribution and returns a random variable
    that is bounded (non-inclusively) by min and max.
"""
function genTruncated(min,max,average,sigma,normalDistribution)
    x=0.0
    while (x < min) || (x > max)
        x=rand(normalDistribution,1)[1]*sigma+average
    end
    return x
end

"""
    Given the eigenvalues λ return a
    vector of random τ, constrained by the λ
"""
function getτ(λ1,λ2,λ3)
    Maxτ1 = 1-abs(λ1)
    Maxτ2 = 1-abs(λ2)
    Maxτ3 = 1-abs(λ3)
    τ1 = (Maxτ1)*(2*rand()-1)
    τ2 = (Maxτ2)*(2*rand()-1)
    τ3 = (Maxτ3)*(2*rand()-1)
    return [τ1 τ2 τ3]
end

"""
    q(l) and Zη(t,l) See equations 20 and 21 of 1707.06926 (Rudnicki et al)
"""
function q(l)
 return (1+l[1]+l[2]+l[3])*(1 + l[1]-l[2]-l[3])*(1-l[1]+l[2]-l[3])*(1-l[1]-l[2]+l[3])
end

"""
    q(l) and Zη(t,l) See equations 20 and 21 of 1707.06926 (Rudnicki et al)
"""
function Zη(t,l)
    norm(t)^4-2*norm(t)^2-2*sum([l[i]^2*(2*(t[i]^2)-norm(t)^2) for i=1:3]) + q(l)
end

"""
    Main function to generate the 'internal' matrix of the channel i.e. the one
    in the form where the unital part consists of diagonal numbers, and the we have
    a non-unital part. This can then be acted on by unitaries either side (rotations)
    to generate an arbitarty CPTP channel
"""
function genSigma(min,max,average,sigma)
    λ1=genTruncated(min,max,average,sigma,_d)
    λ2=genTruncated(min,max,average,sigma,_d)
    λ3=genTruncated(min,max,average,sigma,_d)
    while 1+λ3 < abs(λ1+λ2) || 1-λ3 < abs(λ2-λ1)
        λ3=genTruncated(min,max,average,sigma,_d)
    end
    τ = getτ(λ1,λ2,λ3)
    while norm(τ)^2 > 1-sum([x^2 for x in [λ1 λ2 λ3]])+2*λ1*λ2*λ3 || Zη(τ,[λ1 λ2 λ3]) < 0
        τ = getτ(λ1,λ2,λ3)
    end
    return [1 0 0 0;τ[1] λ1 0 0;τ[2] 0 λ2 0;τ[3] 0 0 λ3]
end

"""
    genChannel(r1,sm,sv,r2)
    Rotate the sigma channel randomly to generate an arbitrary channel
    r1 represents a rotation. The elements of a vector of three normally distributed elements are multiplied by this and used as an XZX rotation. For high fidelity channels this will be small.
    sm, represents the diagonal channel. This is the mean. High fidelity channels will have this close to 1.
    sv, represents the variance in the diagonal channel.
    r2, like r1.

    For example randomFidelityNoise is defined as   genChannel(0.06,0.998,0.04,0.06)
    whereas the noisier randomMeasureNoise is       genChannel(0.05,0.98,0.15,0.05)
"""
function genChannel(r1,sm,sv,r2)
    createRotation(rand(_d,3)*r1) * genSigma(0.9,1,sm,sv) * createRotation(rand(_d,3)*r2)'
end

"""
    given a vector of three angles, generate a random rotation channel
    by using the vector to control a Z*X*Z rotation
"""
function createRotation(angles)
    rotation = Matrix{Float64}(I,4,4)
    rotation[2:4,2:4] =
        [cos(angles[1]) sin(angles[1]) 0;-sin(angles[1]) cos(angles[1]) 0;0 0 1] *
        [1 0 0;0 cos(angles[2]) sin(angles[2]);0 -sin(angles[2]) cos(angles[2])] *
        [cos(angles[3]) sin(angles[3]) 0;-sin(angles[3]) cos(angles[3]) 0;0 0 1]
    return rotation
end

"""
    Helper function to generate typical, high fidelity channel.
"""
function randomFidelityNoise()
    return genChannel(0.06,0.998,0.04,0.06)
end

"""
    Helper function to generate typical, lower fidelity channel.
    I use it to simulate a state preparation channel.
"""
function randomPrepNoise()
    return genChannel(0.05,0.985,0.15,0.05)
end

"""
    Helper function to generate typical, lower fidelity channel.
    I typically use it to simulate measurement error.
"""
function randomMeasureNoise()
    return genChannel(0.05,0.98,0.15,0.05)
end

"""
    function genChannelMap(f)
    Simple helper function - designed to help generate a map of a specific fidelity (supplied parameter).
    Just generates random channels until you get one of the correct fidelity \$\\pm 0.0001\$.
    In general though you would want to rewrite this to give genChannel
    parameters that will result in random channels that often bracket the desired fidelity
    May take a long time if the parameters are far from the types of channels generated
    by randomFidelityNoise.
"""
function genChannelMap(f)
    while true
        candidate = randomFidelityNoise()
        if abs(f-fidelity(candidate)) < 0.0001
            return candidate
        end
    end
end

pI=[1 0;0 1]
pZ=[1 0;0 -1]

# This can be a bit slow, because of the makeSuper on large qubits.
# I have done this better somewhere.
"""
    function add2QPauliNoise(originalMap,cbit,tbit;pauliToAdd = pZ)

    Adds some noise between two qubits in the noise map.
    You pass in an  noise map (super operator basis)
    the control bit and the target bit and optionally the axis of rotation for the noise.
    will return a noise map with a small amount of random noise added between these qubits.
    The noise itself is calculated as \$e^(i*\\frac{\\text{rand()}}{20}\\text{pauliToAdd})\$
"""
function add2QPauliNoise(originalMap,cbit,tbit;pauliToAdd = pZ)
    nqubitsInMap = round(Int,log2(sqrt(size(originalMap,1))))
    #print("Map is $nqubitsInMap qubits\n")
    ok = false
    ## So we are getting irritating (what I think are) rounding errors here, best way to avoid.
    zz = []
    zeroSelect = SparseArrays.sparse([1; 0]*[1 0])
    oneSelect  = SparseArrays.sparse([0;1]*[0 1])
    spI = SparseArrays.sparse(pI)
    thetaa = rand()/10
    thetab = rand()/10
    if cbit < tbit
            startC=cbit
            endC=tbit
             l1 = SparseArrays.sparse(exp(im*thetaa/2*pI))
             l2 = SparseArrays.sparse(exp(im*thetab/2*pauliToAdd))
             for i=cbit+1:tbit-1
                   l1=kron(spI,l1)
                   l2=kron(spI,l2)
             end
            cbitbit = kron(zeroSelect,l1)+kron(oneSelect,l2)
    else
            startC=tbit
            endC=cbit
            l1 = exp(im*thetaa/2*pI)
            l2 = exp(im*thetab/2*pauliToAdd)
            for i=tbit+1:cbit-1
                 l1=kron(l1,spI)
                 l2=kron(l2,spI)
            end
            cbitbit = kron(l1,zeroSelect)+kron(l2,oneSelect)
    end
    for i = 1:(min(cbit,tbit)-1)
            cbitbit = pI ⊗ cbitbit
    end
    for i = max(cbit,tbit)+1:nqubitsInMap
            cbitbit = cbitbit ⊗ pI
    end
    zz = makeSuper(Array(cbitbit))
    withZ = zz*originalMap
    return withZ
end


## Run MMALA using a Bayesian probit model with a Normal prior N(0, priorVar*I)
## on the Swiss banknote data

using Distributions, GeometricMCMC

include("vaso.jl")
include("probitNormalPrior.jl")

# Create Model instance
model = Model(nPars, data, logPrior, logLikelihood, gradLogPosterior,
  tensor, derivTensor, randPrior);

# Create setDriftStep function for adjusting the drift step of MMALA
driftSteps = [1, 1e-8, 1e-7, 1e-6, 1e-5, 1e-4, 1e-3, 1e-2, 1e-1, 0.25]

setDriftStep(i::Int, acceptanceRatio::Float64, nMcmc::Int,
  nBurnin::Int, currentStep::Float64) = setMmalaDriftStep(i, acceptanceRatio,
  nMcmc, nBurnin, currentStep, driftSteps)

# Create instance malaOpts of MMALA options
mmalaOpts = MmalaOpts(55000, 5000, setDriftStep);

# Run MMALA
mcmc = Array(Float64, nPars, nPars)
z = Array(Float64, nPars, nPars)
linearZvMcmc = Array(Float64, nPars, nPars)
linearCoef = Array(Float64, nPars, nPars)
quadraticZvMcmc = Array(Float64, nPars, nPars)
quadraticCoef = Array(Float64, convert(Int, nPars*(nPars+3)/2), nPars)

try
  mcmc, z = mmala(model, mmalaOpts);

  # Compute ZV-MMALA mean estimators based on linear polynomial
  linearZvMcmc, linearCoef = linearZv(mcmc, z);

  # Compute ZV-MMALA mean estimators based on quadratic polynomial
  quadraticZvMcmc, quadraticCoef = quadraticZv(mcmc, z);
catch msg
  println(msg)
end


#=
KTrussTest:
- Julia version: 1.6.2
- Author: samuel
- Date: 2021-09-11
=#

using SuiteSparseGraphBLAS
using SparseArrays
using LinearAlgebra
using Test

include("/Users/samuel/Documents/Jupyter/GraphBlasAlgorithmsFinal/src/KTruss.jl")
matrix =  GBMatrix([[0, 1, 1, 0, 0] [1, 0, 1, 1, 1] [1, 1, 0, 0, 1] [0, 1, 0, 0, 1] [0, 1, 1, 1, 0]])
@test ktruss(matrix, 3, 5) == [[nothing, 1, 1, nothing, nothing]  [1, nothing, 2, 1, 2] [1, 2, nothing, nothing, 1] [nothing, 1, nothing, nothing, 1] [nothing, 2, 1, 1, nothing]]


using Random, LinearAlgebra

# Standardised tests for the eigen decomposition of a square kronecker product
function eigen_tests(rng, C::AbstractKroneckerProduct)

    # Check approximate correctness of decomposition
    λ, Γ = eigen(C)
    @test Γ * Diagonal(λ) * Γ' ≈ C

    # Check approximate correctness of shifted decomposition
    σ² = abs2(randn(rng))
    λ′, Γ′ = (eigen(C) + σ² * I)
    @test Γ′ * Diagonal(λ′) * Γ′' ≈ Matrix(C + σ² * I)
    @test σ² * I + C == C + σ² * I

    # Test with decomposition
    v = randn(rng, length(λ))
    @test eigen(C) \ v ≈ Float64.(Matrix(C)) \ v

    # Test various linear algebra operations with decomposition
    C_dense = Float64.(Matrix(C))
    @test det(eigen(C)) ≈ det(C_dense)
    @test logdet(eigen(C)) ≈ logdet(C_dense)
    @test Matrix(inv(eigen(C))) ≈ inv(C_dense)
end

@testset "eigen" begin
    rng, P, Q, R = MersenneTwister(123456), 3, 5, 7

    # Generate some positive definite matrices so that logdet can be tested.
    A_, B_, C_ = randn(rng, P, P), randn(rng, Q, Q), randn(rng, R, R)
    A, B, C = Symmetric(A_ * A_' + I), Symmetric(B_ * B_' + I), Symmetric(C_ * C_' + I)

    D = kronecker(A, B)
    eigen_tests(rng, D)
    eigen_tests(rng, kronecker(D, C))
    eigen_tests(rng, kronecker(C, D))
    eigen_tests(rng, kronecker(kronecker(A, B), kronecker(B, A)))
    eigen_tests(rng, kronecker(A, 4))
end


using DelimitedFiles

# returns solution of part 1
function x()
	t = vec(readdlm("10.txt", Int64))
	sort!(t)

	a=b=0
	prev = 0
	for ch in t
		if ch - prev == 1
			a += 1
		elseif ch - prev == 3
			b += 1
		end
		prev = ch
	end
	b += 1
	a*b
end

# returns solution of part 2
function y()
	t = sort(vec(readdlm("10.txt", Int64)))

	nums = Array{Int64,1}(zeros(maximum(t)+3))

	for ch in t
		s = ch <= 3 ? 1 : 0
		for i = max(1, ch-3):ch-1
			s += nums[i]
		end
		nums[ch] = s
	end
	nums[maximum(t)]
end


# function related to calculating areas and volumes

@doc """
### calcVolumeContribution!

Returns the contribution from the given faces to the volume of enclosed region.
The Gauss-Divergence theorem is applied to (1/3) ∫∫∫ ∇⋅(x,y,z) dΩ = vol.

**Inputs**

* `sbpface`: an SBP face operator type
* `xsbp`: SBP-face nodes in physical space; [coord, sbp node, face]
* `nrm`: scaled face-normal at the sbpface nodes; [component, sbp node, face]
* `vol`: volume contribution from boundaries other than this set of faces

**Returns**

* `vol`: the updated value for the volume

"""->
function calcVolumeContribution!(sbpface::AbstractFace{Tsbp},
                                 xsbp::AbstractArray{Tmsh,3},
                                 nrm::AbstractArray{Tmsh,3}) where {Tsbp,Tmsh}
  @assert( sbpface.numnodes == size(xsbp,2) == size(nrm,2) )
  @assert( size(xsbp,3) == size(nrm,3) )
  @assert( size(xsbp,1) == size(nrm,1) )
  vol = zero(Tmsh)
  for f = 1:size(nrm,3)
    for i = 1:sbpface.numnodes
      for di = 1:size(nrm,1)# 3
        vol += (1/3)*sbpface.wface[i]*xsbp[di,i,f]*nrm[di,i,f]
      end
    end
  end
  return vol
end

"""
  This function computes the volume contribution of a set of faces, where the
  set of faces is given by which boundary condition they have applied to them.
"""
function calcVolumeContribution!(mesh::AbstractMesh{Tmsh}, eqn::AbstractSolutionData, bndry_nums::Array{Int, 1}) where Tmsh
# compute the volume contribution for the faces associated with a set of
# boundary conditions

  vol = zero(Tmsh)
  for i=1:length(bndry_nums)
    start_idx = mesh.bndry_offsets[ bndry_nums[i] ]
    end_idx = mesh.bndry_offsets[ bndry_nums[i] + 1 ] - 1
    face_range = start_idx:end_idx
    bndry_faces = sview(mesh.bndryfaces, face_range)

    # calculate normal
    nrm = sview(mesh.nrm_bndry, :, :, face_range)
#    nrm = computeNormal(mesh, eqn, bndry_faces, mesh.dxidx_bndry)
    coords = sview(mesh.coords_bndry, :, :, face_range)

    # compute volume contribution
    vol += calcVolumeContribution!(mesh.sbpface, coords, nrm)
  end

  return vol
end

function calcVolumeContribution_rev!(mesh::AbstractMesh{Tmsh}, eqn::AbstractSolutionData, bndry_nums::Array{Int, 1}, vol_bar::Tmsh) where Tmsh

  vol = zero(Tmsh)
  for i=1:length(bndry_nums)
    start_idx = mesh.bndry_offsets[ bndry_nums[i] ]
    end_idx = mesh.bndry_offsets[ bndry_nums[i] + 1 ] - 1
    face_range = start_idx:end_idx
    bndry_faces = sview(mesh.bndryfaces, face_range)
    coords = sview(mesh.coords_bndry, :, :, face_range)
    coords_bar = zeros(coords)  #TODO: make this a field of mesh?

    nrm = sview(mesh.nrm_bndry, :, :, face_range)
    nrm_bar = sview(mesh.nrm_bndry_bar, :, :, face_range)
#    nrm = computeNormal(mesh, eqn, bndry_faces)
#    nrm_bar = zeros(Tmsh, mesh.dim, mesh.numNodesPerFace, length(bndry_faces))

    # compute volume contribution
    calcVolumeContribution_rev!(mesh.sbpface, coords, coords_bar, nrm, nrm_bar, vol_bar)

    # calculate normal
#    computeNormal_rev(mesh, eqn,  bndry_faces, nrm_bar)

  end

  return nothing
end

"""
  Computes the normal vector at each surface node of the given list of faces.
  Straight-sided meshes only!

  Inputs:
    mesh: an AbstractMesh
    sbp: an SBP operator
    bndryfaces: list of boundary faces

  Outputs: nrm: mesh.dim x mesh.numNodesPerFace x length(bndryfaces) array
                containing the normal vectors
"""
function computeNormal(mesh::AbstractMesh{Tmsh}, eqn::AbstractSolutionData,
                       bndryfaces::AbstractArray{Boundary, 1},
                       dxidx::AbstractArray{Tmsh, 4}) where Tmsh

  if typeof(mesh) <: AbstractDGMesh
    @assert mesh.coord_order == 1
  end

  nfaces = length(bndryfaces)
  nrm = zeros(Tmsh, mesh.dim, mesh.numNodesPerFace, nfaces)

  for i=1:nfaces
    bndry_i = bndryfaces[i]
    nrm_xi = sview(mesh.sbpface.normal, :, bndry_i.face)
    for j=1:mesh.numNodesPerFace
      dxidx_j = sview(dxidx, :, :, j, i)
      nrm_xy_j = sview(nrm, :, j, i)
      calcBCNormal(eqn.params, dxidx_j, nrm_xi, nrm_xy_j)
    end
  end

  return nrm
end

"""
  This function uses reverse mode to back propigate perturbations in nrm to
  dxidx_bndry
"""
function computeNormal_rev(mesh::AbstractMesh, eqn::AbstractSolutionData,
                           bndryfaces::AbstractArray{Boundary, 1},
                           nrm_bar::Abstract3DArray{Tmsh},
                           dxidx::AbstractArray{Tmsh, 4},
                           dxidx_bar::AbstractArray{Tmsh, 4}) where Tmsh
  if typeof(mesh) <: AbstractDGMesh
    @assert mesh.coord_order == 1
  end

  nfaces = length(bndryfaces)
  for i=1:nfaces
    bndry_i = bndryfaces[i]
    nrm_xi = sview(mesh.sbpface.normal, :, bndry_i.face)
    for j=1:mesh.numNodesPerFace
      dxidx_j = sview(dxidx, :, :, j, i)
      dxidxbar_j = sview(dxidx_bar, :, :, j, i)
      nrmbar_xy_j = sview(nrm_bar, :, j, i)
      calcBCNormal_revm(eqn.params, dxidx_j, nrm_xi, nrmbar_xy_j, dxidxbar_j)
    end
  end

  return nothing
end

@doc """
### calcVolumeContribution_rev!

This is the reverse differentiated version of calcVolumeContribution!.  See docs
of calcVolumeContribution! for further details of the primal method.  This
function is differentiated with respect to the primal version's `xsbp` and `nrm`
variables.

**Inputs**

* `sbpface`: an SBP face operator type
* `xsbp`: SBP-face nodes in physical space; [coord, sbp node, face]
* `nrm`: scaled face-normal at the sbpface nodes; [component, sbp node, face]
* `vol_bar`: left multiplies d(vol)/d(nrm) and d(vol)/d(xsbp); it is a scalar

**InOuts**

* `xsbp_bar`: result of vector Jacobian product; [component, sbp node, face]
* `nrm_bar`: result of vector Jacobian product; [component, sbp node, face]

"""->
function calcVolumeContribution_rev!(sbpface::AbstractFace{Tsbp},
                                     xsbp::AbstractArray{Tmsh,3},
                                     xsbp_bar::AbstractArray{Tmsh,3},
                                     nrm::AbstractArray{Tmsh,3},
                                     nrm_bar::AbstractArray{Tmsh,3},
                                     vol_bar::Tmsh) where {Tsbp,Tmsh}
  @assert( sbpface.numnodes == size(xsbp,2) == size(xsbp_bar,2) == size(nrm,2)
           == size(nrm_bar,2) )
  @assert( size(xsbp,3) == size(xsbp_bar,3) == size(nrm,3) == size(nrm_bar,3) )
  @assert( size(xsbp,1) == size(xsbp,1) == size(nrm,1) == size(nrm_bar,1) )
  for f = 1:size(nrm,3)
    for i = 1:sbpface.numnodes
      for di = 1:size(nrm,1) # 3
        # vol += (1/3)*sbpface.wface[i]*xsbp[di,i,f]*nrm[di,i,f]
        fac = vol_bar*(1/3)*sbpface.wface[i]
        xsbp_bar[di,i,f] += fac*nrm[di,i,f]
        nrm_bar[di,i,f] += fac*xsbp[di,i,f]
      end
    end
  end
end

@doc """
### calcProjectedAreaContribution!

  Returns the contribution from the given faces to the projected area onto the
  plane of coordiante `di`

  Because this function uses the absolute value of the normal vector, it is not
  differentiable at zero.

**Inputs**

* `sbpface`: an SBP face operator type
* `di`: the index that indicates the projection plane (e.g. `di`=1 --> yz plane)
* `nrm`: scaled face-normal at the sbpface nodes; [component, sbp node, face]
* `projarea`: projected area of faces from other boundaries.

**Returns**

* `projarea`: the updated value for projected area

"""->
function calcProjectedAreaContribution!(sbpface::AbstractFace{Tsbp},
                                        nrm::AbstractArray{Tmsh,3},
                                        di::Int) where {Tsbp,Tmsh}
  @assert( sbpface.numnodes == size(nrm,2) )
  projarea = zero(Tmsh)
  for f = 1:size(nrm,3)
    for i = 1:sbpface.numnodes
      projarea += 0.5*sbpface.wface[i]*abs(nrm[di,i,f])
    end
  end
  return projarea
end

"""
  This function computes the volume contribution of a set of faces, where the
  set of faces is given by which boundary condition they have applied to them.
"""
function calcProjectedAreaContribution!(mesh::AbstractMesh{Tmsh}, eqn::AbstractSolutionData, bndry_nums::Array{Int, 1}, di::Int) where Tmsh

  proj_area = zero(Tmsh)
  for i=1:length(bndry_nums)
    start_idx = mesh.bndry_offsets[ bndry_nums[i] ]
    end_idx = mesh.bndry_offsets[ bndry_nums[i] + 1 ] - 1
    face_range = start_idx:end_idx
    bndry_faces = sview(mesh.bndryfaces, face_range)

    # calculate normal
    nrm = sview(mesh.nrm_bndry, :, :, face_range)
#    nrm = computeNormal(mesh, eqn, bndry_faces)
#    coords = sview(mesh.coords_bndry, :, :, face_range)

    proj_area += calcProjectedAreaContribution!(mesh.sbpface, nrm, di)
  end

  return proj_area
end

@doc """
### calcProjectedAreaContribution_rev!

This is the reverse differentiated version of calcProjectedAreaContribution!.
See docs of calcProjectedAreaContribution! for further details of the primal
method.  This function is differentiated with respect to the primal version's
`nrm` variable.

**Inputs**

* `sbpface`: an SBP face operator type
* `di`: the index that indicates the projection plane (e.g. `di`=1 --> yz plane)
* `nrm`: scaled face-normal at the sbpface nodes; [component, sbp node, face]
* `projarea_bar`: left multiplies d(projarea)/d(nrm); it is a scalar

**InOuts**

* `nrm_bar`: result of vector Jacobian product; [component, sbp node, face]

"""->
function calcProjectedAreaContribution_rev!(sbpface::AbstractFace{Tsbp}, nrm::AbstractArray{Tmsh,3},
    nrm_bar::AbstractArray{Tmsh,3}, di::Int, projarea_bar::Tmsh) where {Tsbp,Tmsh}
  @assert( sbpface.numnodes == size(nrm,2) == size(nrm_bar,2) )
  for f = 1:size(nrm,3)
    for i = 1:sbpface.numnodes
      # projarea += 0.5*sbpface.wface[i]*abs(nrm[di,i,f])
      #  v1 = nrm[di,i,f]
      #  v2 = abs(v1)
      #  v3 = 0.5*sbpface.wface[i]
      #  v4 = v2*v3

      nrm_bar[di,i,f] += (projarea_bar*0.5*sbpface.wface[i]*
                          absvalue_deriv(nrm[di,i,f]))
    end
  end
end

"""
  Back propigate projarea_bar to mesh.dxidx_bndry_bar
"""
function calcProjectedAreaContribution_rev!(mesh::AbstractMesh, eqn::AbstractSolutionData, bndry_nums::Array{Int, 1}, di::Int,  projarea_bar::Tmsh) where Tmsh

  for i=1:length(bndry_nums)
    start_idx = mesh.bndry_offsets[ bndry_nums[i] ]
    end_idx = mesh.bndry_offsets[ bndry_nums[i] + 1 ] - 1
    face_range = start_idx:end_idx
    bndry_faces = sview(mesh.bndryfaces, face_range)

    nrm = sview(mesh.nrm_bndry, :, :, face_range)
    nrm_bar = sview(mesh.nrm_bndry_bar, :, :, face_range)
#    nrm = computeNormal(mesh, eqn, bndry_faces)
#    nrm_bar = zero(nrm)
    calcProjectedAreaContribution_rev!(mesh.sbpface, nrm, nrm_bar, di, projarea_bar)

    # back propigate to dxidx_bndry_bar
#    computeNormal_rev(mesh, eqn, bndry_faces, nrm_bar)
  end

  return nothing
end

function crossProd(x::AbstractArray{T,1}, y::AbstractArray{T,1},
                   cp::AbstractArray{T,1}) where T
  cp[1] = x[2]*y[3] - y[2]*x[3]
  cp[2] = x[3]*y[1] - y[3]*x[1]
  cp[3] = x[1]*y[2] - y[1]*x[2]
end

function crossProd_rev(x::AbstractArray{T,1}, x_bar::AbstractArray{T,1},
                       y::AbstractArray{T,1}, y_bar::AbstractArray{T,1},
                       cp_bar::AbstractArray{T,1}) where T
  #cp[1] = x[2]*y[3] - y[2]*x[3]
  x_bar[2] +=  y[3]*cp_bar[1]
  y_bar[3] +=  x[2]*cp_bar[1]
  y_bar[2] += -x[3]*cp_bar[1]
  x_bar[3] += -y[2]*cp_bar[1]
  #cp[2] = x[3]*y[1] - y[3]*x[1]
  x_bar[3] +=  y[1]*cp_bar[2]
  y_bar[1] +=  x[3]*cp_bar[2]
  y_bar[3] += -x[1]*cp_bar[2]
  x_bar[1] += -y[3]*cp_bar[2]
  #cp[3] = x[1]*y[2] - y[1]*x[2]
  x_bar[1] +=  y[2]*cp_bar[3]
  y_bar[2] +=  x[1]*cp_bar[3]
  y_bar[1] += -x[2]*cp_bar[3]
  x_bar[2] += -y[1]*cp_bar[3]
end


include("./exportData.jl")
using .exportData, DifferentialEquations, DiffEqParamEstim, PyPlot, Optim, DelimitedFiles

# define fenkar system
function H(x;k=100.0)
	return 0.5*(1.0 + tanh(k*x))
end

function fenkar!(du, u, p, t)

	# parameters
	@views tsi,tv1m,tv2m,tvp,twm,twp,td,to,tr,xk,uc,uv,ucsi,t0,TI,IA = p[1:16]

	# fenkar dynamics
	du[1] = H(t-t0;k=1.0)*IA*sin(pi*(t-t0)/TI)^500 - (u[1]*H(uc-u[1])/to + H(u[1]-uc)/tr - u[2]*H(u[1]-uc)*(1.0-u[1])*(u[1]-uc)/td - u[3]*H(u[1]-ucsi;k=xk)/tsi)
	du[2] = H(uc-u[1])*(1.0-u[2])/(tv1m*H(uv-u[1]) + tv2m*H(u[1]-uv)) - H(u[1]-uc)*u[2]/tvp
	du[3] = H(uc-u[1])*(1.0-u[3])/twm - H(u[1]-uc)*u[3]/twp

	return nothing
end

# fenkar system parameters in p[1:13] (BR; from https://doi.org/10.1063/1.166311)
# and then p[14:16] = [offset_time, stimulus_period, stimulus_amplitude]
p = [29.0,19.6,1250.0,3.33,41.0,870.0,0.25,12.5,33.3,10.0,0.13,0.04,0.85,0.0,125.0,0.05]
tspan = (0.0, 3000.0)
u0 = zeros(Float64,3)

#solve and take data
prob = ODEProblem(fenkar!, u0, tspan, p)
sol = solve(prob, Tsit5(); saveat=tspan[begin]:2.0:tspan[end])

# getExpData
ind = 114
t, data, p[15] = getExpData(ind, tInds=1:1000)
p[14] = t[findfirst(data[1,:] .> 0.5)]-p[15]/2

# remake problem
prob = remake(prob, tspan=(0.0, t[end]), p=p)

# use the save_idxs to select the mid-point of the u-variable
sol = solve(prob, Tsit5(); saveat=t, save_idxs=1:1)

fig = plt.figure(figsize=(4,2))
plt.plot(t, data[:], "-", label="Exp. Sample")
plt.plot(sol.t, sol[1,:], "--", label="Initial Parameters")
plt.xlabel("\$ t \$")

# define the cost function
cost_function = build_loss_objective(prob, Tsit5(), L2Loss(t,data), maxiters=10000, verbose_opt=true, verbose_steps=50; save_idxs=1:1);

# define limits for parameters
#          tsi,  tv1m,   tv2m,   tvp,   twm,    twp,  td,    to,    tr,   xk,  uc,    uv, ucsi,        t0,     TI,   IA
#   p = [ 29.0,  19.6, 1250.0,   3.33, 41.0,  870.0, 0.25, 12.5,  33.3, 10.0, 0.13, 0.04,  0.85,      1.0 ,   BCL, 0.05]
lower = [  1.0,  10.0,  400.0,   2.33, 10.0,  100.0, 0.10, 10.0,  10.0,  9.0, 0.10, 0.01,  0.10,      0.0,   40.0, 0.00]
upper = [ 50.0,  50.0, 2000.0,   4.33, 50.0, 1000.0, 0.50, 50.0,  50.0, 11.0, 0.90, 0.90,  0.90,     t[1], 1000.0, 1.00]

@assert all(p .>= lower) && all(p .<= upper)

# fit
result = optimize(cost_function, lower, upper, p, Fminbox(NelderMead()))
@info "Initial p = $(p) \n"
@info "Minimizer = $(result.minimizer) \n"
prob = remake(prob, p=result.minimizer)
sol = solve(prob, Tsit5(); saveat=t, save_idxs=1:1)
plt.plot(sol.t, sol[1,:], "--", label="Nelder-Mead Parameters")

result = optimize(cost_function, lower, upper, result.minimizer, Fminbox(BFGS()))
@info "Initial p = $(p) \n"
@info "Minimizer = $(result.minimizer) \n"
prob = remake(prob, p=result.minimizer)
sol = solve(prob, Tsit5(); saveat=t, save_idxs=1:1)
plt.plot(sol.t, sol[1,:], "--", label="BFGS Parameters")

plt.savefig("./ode/optim/$(ind)_fittings.pdf")
plt.close()

open("./ode/optim/$(ind)_params.txt", "w") do io
	writedlm(io, ["tsi" "tv1m" "tv2m" "tvp" "twm" "twp" "td" "to" "tr" "xk" "uc" "uv" "ucsi" "t0" "TI" "IA"])
	writedlm(io, "\n")
	writedlm(io, transpose(p))
	writedlm(io, "\n")
	writedlm(io, transpose(result.minimizer))
end


using GLPKMathProgInterface, Ipopt, ECOS

lp_solver = GLPKSolverLP()
ip_solver = GLPKSolverMIP()
conic_solver = ECOSSolver(verbose=false)
nlp_solver = IpoptSolver(print_level=0, fixed_variable_treatment="make_constraint")

include("linprog.jl")
linprogtest(lp_solver)

include("mixintprog.jl")
mixintprogtest(ip_solver)

include("quadprog.jl")
quadprogtest(nlp_solver)
qpdualtest(nlp_solver)
socptest(conic_solver)

include("nlp.jl")
nlptest(nlp_solver)
nlptest_nohessian(nlp_solver)
convexnlptest(nlp_solver)
rosenbrocktest(nlp_solver)

include("conicinterface.jl")
coniclineartest(conic_solver, duals=true)
# Test conic fallback for LPs
coniclineartest(lp_solver, duals=true)

include("linproginterface.jl")
linprogsolvertest(lp_solver)
linprogsolvertestextra(lp_solver)
# Test LP fallback for conics
linprogsolvertest(conic_solver)
# Test LP fallback for nlp
linprogsolvertest(nlp_solver, 1e-5)


using LinearAlgebra
using Random, Distributions
using Plots
Plots.scalefontsizes(1.5)

# Global Terms
Random.seed!(1234)

# Logistic Regression

function f(β)
    #=
    Objective Function
    x: n x d
    y: d x 1
    β: d x 1
    λ: penalty
    =#
    result =  sum((X * β - y) .^ 2) / 2 / n + λ *norm(β,1)
    return result
end

function ∇f(β)
    #=
    This is already noisy gradient.
    β: 100 x 1
    σ: noise
    =#
    return X' * (X * β - y) ./n .+ λ .* sign.(β) .+ rand(Normal(0,σ^2), (d,1))
end

# MASG

function MASG(f, ∇f, nk, α, μ)
    #=
        f:  objective functin
        ∇f: gradient of f
        X:  design matrix
        y:  {-1,1} class
        a_n:an array containing each step each iteration
        K:  total number of iterations
        α:  an array containing step sizes
        μ:  μ-strongly convex
        λ:  penalty
    =#
    x_0 = x_1 = zeros(d)

    # Record Histories
    f_history = zeros(sum(nk))
    γ = 1

    for k in 1:K
        x_0 = x_1
        for m in 1:nk[k]
            w = (1 - sqrt(μ * α[k])) / (1 + sqrt(μ * α[k]))
            ymk = (1 + w) .* x_1 - w .* x_0
            x_0 = x_1
            x_1 = ymk .- α[k] .* ∇f(ymk)

            # Record History
            f_history[γ] = f(x_1)
            γ = γ + 1
        end
    end
    return f_history
end

function compute_nk(κ, p, Δ, label)
    C = floor(sqrt(κ) * log(2^(p+2))) + 1
    grid = [1:K;]
    nk = 2 .^ grid
    if label == 3.7
        nk[1] = floor(sqrt(κ) * log(κ)) + 1
    elseif label == 4
        nk[1] = floor(sqrt(κ) * log(2 * L * Δ / (σ^2 * sqrt(κ)))) + 1
    end
    return nk
end

function compute_MASG_α(L)
    grid = [1:K;]
    α = 1 ./ (2 .^ (2 .* [1:K;]) .* L)
    α[1] = 1 / L
    return α
end

# Vanilla Gradient Descent
function GD(f, ∇f)
    x = zeros(d)
    f_history = zeros(η)
    for i in 1:η
        x = x - 1 / L * ∇f(x)
        f_history[i] = f(x)
    end
    return f_history
end

# Accelerated Gradient Descent
function AGD(f, ∇f)
    x_0 = x_1 = zeros(d)
    f_history = zeros(η)
    for i in 1:η
        y = x_1 + (i-1)/(i+2) *(x_1-x_0)
        x_0 = x_1
        x_1 = y - 1/L * ∇f(y)

        f_history[i] = f(x_1)
    end
    return f_history
end

# AC-SA
## Bregman Distance
function D(β1, β2)
    #=
    Prox function, Bregman Distance
    =#
    return f(β1) - f(β2) - (∇f(β2)' * (β1 .- β2))[1]
end

function ORDA(τ)
    #=
    ξ is the turning param
    M is constant
    τ is quadratic constant
    =#
    M = 0
    c = (sqrt(τ) * (σ + M))/(2 * sqrt(D(βstart, zeros(d))) )
    θ(t) = 2/(t+2)
    ν(t) = 2/(t+1)
    γ(t) = c * (t+1)^(3/2)+τ*L
    G = 0
    zt = z0 = xt = x0 = zeros(d)
    f_history = zeros(η+1)
    for t in 0:η
        f_history[t+1] = f(xt)
        p1 = (1 - θ(t)) * (μ + θ(t)^2 * γ(t)) / (θ(t)^2 * γ(t) + (1-θ(t)^2) * μ) .* xt
        p2 = ((1 - θ(t)) * θ(t) * μ + θ(t)^3 * γ(t)) / (θ(t)^2 * γ(t) + (1-θ(t)^2) * μ) * zt
        yt = p1 + p2
        G = G .+ ∇f(yt) ./ ν(t)
        gt = θ(t) * ν(t) .* G

        zp1 = ∇f(x0) .- λ .- (gt .+ λ) ./ (θ(t) * ν(t) * γ(t+1))
        zt = inv(X' * X) * (zp1 .* n + X' * y)

        xp1 = ∇f(yt)  .- λ .- (∇f(yt) .+ λ) / (μ / (τ * θ(t)^2) + γ(t) / τ)
        xt = inv(X' * X) * (xp1 .* n + X' * y)
     end
     return f_history
end

function ACSA(ν)
    α(t) = 2/(t+1)
    γ(t) = 4 * 2 * L / (ν * t * (t+1))

    xt = xtag = zeros(d)
    f_history = zeros(η+1)
    f_history[1] = f(xt)

    for t in 1:η
        xp1 = (1 - α(t)) * (μ + γ(t)) / (γ(t) + (1-α(t)^2) * μ) * xtag
        xp2 = α(t)*((1-α(t)) * μ +γ(t))/(γ(t) + (1-α(t)^2) * μ) * xt
        xtmd = xp1 + xp2

        p1 = (∇f(xtmd) * (μ-1) * α(t) .- α(t) * λ .+ ((1-α(t))*μ + γ(t))*∇f(xt)) / (μ + γ(t))
        p2 = (p1 .- λ) .*n + X' * y
        xt = inv(X' * X) * p2

        xtag = α(t) * xt + (1- α(t)) * xtag
        f_history[t+1] = f(xt)
    end

    return f_history
end

## Data setup
global n = 10000   # number of datapoints
global d = 100    # number of dimensions
global λ = 0.01   # Penalty for beta term
global ρ = 0
global X = rand(Normal(0,1), (n,d))
β = rand(Normal(0, 1), (d, 1))
global βstart = β
global y = X * β  .+ rand(Normal(0,1),(n,1))

# Experiments
global K = 10     # number of iterations for MASG
global η = 2000   # Number of iterations for other methods
global L = maximum(eigvals(X' * X))
# global L = norm(X,2)^2/n
global μ = 0

global σ = sqrt(1e-1)
# MASG
α_MASG = compute_MASG_α(L)
nk_7 = compute_nk(κ, 1, norm(β,2),3.7)
f_MASG_7 = MASG(f, ∇f, nk_7, α_MASG, μ)

## MASG Thm 4
nk_4 = compute_nk(κ, 1, norm(β,2), 4.0)
f_MASG_4 = MASG(f, ∇f, nk_4, α_MASG, μ)

# Vanilla Gradient Descent
f_GD = GD(f, ∇f)

# Vanilla Accelerated Gradient Descent
f_AGD = AGD(f, ∇f)

# ORDA
f_ORDA = ORDA(L)

ν = norm(X,2) / n
f_ACSA = ACSA(ν)

plot(f_GD, label="GD",legend=:bottomleft,w=2,linestyle=:dash,color=:red, xscale = :log, yscale=:log)
plot!(f_AGD, label="AGD", w=2, linestyle=:dash,color=:purple)
plot!(f_ORDA, label = "ORDA", w=2, color=:blue)
plot!(f_ACSA, label="ACSA",w=2, color=:brown)
plot!(f_MASG_7, label="MASG",w=2, linestyle=:solid, color=:green)
plot!(f_MASG_4, label="MASG*",w=2, linestyle=:solid, color=:black)
xlabel!("Iteration")

savefig("./figures/non_sc_10000_4.png")


module KnapsackLib

export KnapData, KnapItem, solveKnapModel, solveKnapNaive, solveKnapExpCore, solveKnapHeur, genKnap

using Random, JuMP

const EPS = 1e-7

include("Data.jl")
include("Model.jl")
include("Naive.jl")
include("ExpandingCore.jl")
include("Heuristic.jl")
include("Generator.jl")

end


using Test
using SimpleGraphs, SimpleGraphAlgorithms, Polynomials

use_Cbc()
G = Paley(17)

# Cliques and independent sets
A = max_clique(G)
B = max_indep_set(G)
@test length(A) == length(B)

# Isomorphism
f = iso(G,G')
@test length(f) == 17
@test is_iso(G,G')

# Edge matching
M = max_matching(G)
@test length(M) == 8

H = Dodecahedron()
X = kfactor(H,3)
@test length(X) == 30

# Vertex coloring
f = color(G,6)
@test length(f) == 17

# Edge coloring
@test edge_chromatic_number(G) == 9
f = edge_color(G,9)
@test length(f) == NE(G)

# Domination
A = min_dom_set(G)
@test length(A) == 3

# Covering
A = min_vertex_cover(G)
@test length(A) == 14
A = min_vertex_cover(G,10)
@test length(A) == 2
A = min_edge_cover(G)
@test length(A) == 9

# MAD
@test mad(G) == maximum(deg(G))

# Chromatic Polynomial
f = chrome_poly(Cycle(5))
@test f(2) == 0
@test f(3) == 30

# Chromatic number
@test chromatic_number(Petersen()) == 3

G = Spindle()
add!(G,4,0)
H = Spindle()
add!(H,7,0)
@test fast_iso_test(G,H)
f = iso(G,H)
@test f[0] == 0

# edge and vertex connectivity
G = Complete(8,8)'
add!(G,1,9)
add!(G,1,10)
@test edge_connectivity(G) == 2
@test connectivity(G)==1


module psfsize

using ..params
using ..psf
using FFTW
export calcsize

function calcsize(p::ParameterSet, obj::Space)
    # we set an arbitrarily high length to test across
    testrange = collect((0:1:(p.sim.subvpix*20)) .* p.sim.subpixelpitch)
    x3max = maximum(abs.(obj.z))

    if p.opt.a0 > zero(Float64)
        steps = 10
        stepz = p.opt.a0/steps
        lineproj = zeros(ComplexF64, length(testrange), obj.zlen)
        linemag = zeros(length(testrange), obj.zlen)
        Hline = fresnelH(lineproj[:,1], p, stepz)
        plan = plan_fft!(lineproj[:,1])
        for i in 1:obj.zlen
            lineproj[:,i] = psfline(testrange, obj.z[i], p)
            plan*lineproj[:,i]
            lineproj[:,i] = lineproj[:,i].*(Hline.^steps)
            plan\lineproj[:,i]
            linemag[:,i] = abs2.(lineproj[:,i]) ./ maximum(abs2.(lineproj[:,i]))
        end
        imgspace = makeimgspace(dropdims(reduce(max, linemag, dims=2); dims=2), p)
    else
        lineproj = psfline(testrange, x3max, p)
        linemag =  abs2.(lineproj) ./ maximum(abs2.(lineproj))
        imgspace = makeimgspace(linemag, p)
    end
    return imgspace
end

function psfline(x₁line::Vector{Float64}, x₃::Float64, p::ParameterSet)

    l²norm² = abs.(x₁line) ./ p.opt.M
    v = l²norm² .* (p.con.k*sin(p.con.alpha))
    u = 4*p.con.k * x₃ * (sin(p.con.alpha/2)^2)
    Kₒ = p.opt.M/(p.opt.fobj*p.opt.lambda)^2 * exp(-im*u/(4(sin(p.con.alpha/2)^2)))

    psfline = Array{ComplexF64,1}(undef,length(v))

    Threads.@threads for i in 1:length(v)
        @inbounds psfline[i] = integratePSF(v[i], u, p.opt.a0, p.con.alpha) * Kₒ
    end

    return psfline
end

function makeimgspace(psfline::Vector{Float64}, p::ParameterSet)
    outArea = psfline .< 0.1
    if sum(outArea) == 0
        error("Estimated PSF size exceeds the limit")
    end
    sizeref = cld(findfirst(outArea)[1],p.sim.subvpix)
    halfwidth = max( p.sim.subvpix*(sizeref + 4), 2*p.sim.subvpix)

    "Create X-Y image space based on halfwidth."
    x =(-halfwidth:1:halfwidth) * p.sim.subpixelpitch
    y =(-halfwidth:1:halfwidth) * p.sim.subpixelpitch
    imgspace = Space(x,y)

    return imgspace
end

end # module


using PyPlot, Images
using PyCall
@pyimport numpy

function DIC_cosine_transform(img_name,lambda)

    # import image
    raw_img = PyPlot.imread(img_name);
    img = convert(Image{Images.Gray},raw_img[:,:,1:3]);
    M = convert(Array,img);
    M = float(M);

    #subtract average of each column
    mx, my = size(M)
    m1 = mean(M[:,1])
    mL = mean(M[:,end])
    for i=1:mx
        M[:,i] = M[:,i] - mean(M[:,i])
    end

    gX, gY = numpy.gradient(M);
    d1 = FFTW.r2r(gY,FFTW.REDFT00)
    nx, ny = size(d1)
    d2=zeros(ny,nx)
    d2=[(nx/pi)^2 * d1[j,i]/(lambda*j^(2.0) + i^2) for i in 1:ny, j in 1:nx];
    #d2=[(nx/pi)^2 * d1[j,i] for i in 1:ny, j in 1:nx];

    dst = convert(Array{Float64,2},d2)
    e1 = FFTW.r2r(dst,FFTW.REDFT00)
    e1 = e1 / (2*nx)^2

    maxi = maximum(e1)
    mini = minimum(e1)
    clamped = 1 +  (-e1 + maxi) * 1 / ( mini-maxi);
    out_name = split(image_name,".")[1] * "_cos_l=$(@sprintf("%.2f", lambda))" * ".png"
    imwrite(grayim(clamped),out_name)

    return 0

end


function calc_Etot( pw::PWGrid, Vpot, psi::Array{ComplexF64,2} )

    Ω = pw.Ω
    Npoints = pw.Npoints
    Nstates = size(psi)[2]

    Kpsi = op_K( pw, psi )
    Ekin = 0.0
    for is = 1:Nstates
        Ekin = Ekin + real( dot( psi[:,is], Kpsi[:,is] ) )
    end

    # Calculate in real space
    rho = calc_rho( pw, psi )
    Epot = dot( rho, Vpot ) * Ω/Npoints

    Etot = Ekin + Epot
    return Etot
end


function c2ir(c::AbstractVector, len=256)
    h = Array{eltype(c),1}(len)
    m = length(c) - 1

    h[1] = exp(c[1])
    for n=2:len
        d = 0.0
        upperlimit = (n >= m+1) ? m+1 : n
        for k=2:upperlimit
            @inbounds d += (k-1) * c[k] * h[n-k+1]
        end
        @inbounds h[n] = d / (n-1)
    end

    h
end


# ODE SIR/SICR models
# variables are concentrations, X/N
# Z = 1-S

function sir(T,N,beta,sigmaR,sigmaD,a,b)

	dI = zeros(Int,T)
	dR = zeros(Int,T)
	dD = zeros(Int,T)
	I = zeros(Int,T)

	S = N
	I[1] = 10
	for i = 2:T
		d = Binomial(S,beta*S*I[i-1]/N^2)
		dI[i] = rand(d)
		d = Binomial(I[i-1],sigmaR*I[i-1]/N)
		dR[i] = rand(d)
		d=Gamma(a,b)
		ind = round(Int,rand(d))
		ind = max(i-1-ind,1)
		d = Binomial(I[i-1],sigmaD*I[ind]/N)
		dD[i] = rand(d)
		I[i] = I[i-1] + dI[i] - dR[i] - dD[i]
		S -=  dI[i]
	end

	return dI, dR, dD

end

function sicr(T,N,beta,sigmaI,sigmaC,sigmaR,sigmaD,a,b)

	dI = zeros(Int,T)
	dC = zeros(Int,T)
	dR = zeros(Int,T)
	dD = zeros(Int,T)
	dC = zeros(Int,T)
	C = zeros(Int,T)
	S = N
	I = 10
	for i = 2:T
		d = Binomial(S,beta*S*I/N^2)
		dIp = rand(d)
		d = Binomial(I,sigmaI*I/N)
		dIm = rand(d)
		d = Binomial(I,sigmaC*I/N)
		dC[i] = rand(d)
		d = Binomial(C[i-1],sigmaR*C[i-1]/N)
		dR[i] = rand(d)
		d=Gamma(a,b)
		ind = round(Int,rand(d))
		ind = max(i-1-ind,1)
		d = Binomial(C[i-1],sigmaD*C[ind]/N)
		dD[i] = rand(d)
		C[i] = C[i-1] + dC[i] - dR[i] - dD[i]
		I += dIp - dIm - dC[i]
		S -=  dIp
	end

	return dC, dR, dD

end

function sicr(T,N,beta,sigmaI,sigmaC,sigmaR,sigmaD)

	dI = zeros(Int,T)
	dC = zeros(Int,T)
	dR = zeros(Int,T)
	dD = zeros(Int,T)
	dC = zeros(Int,T)
	C =1
	R = 0
	S = N
	I = 10
	for i = 2:T
		d = Binomial(S,beta*S*I/N^2)
		dIp = rand(d)
		d = Binomial(I,sigmaI*I/N)
		dIm = rand(d)
		d = Binomial(I,sigmaC*I/N)
		dC[i] = rand(d)
		if C > 0
			d = Binomial(C,sigmaR*C/N)
			dR[i] = rand(d)
		else
			dC[i] = 0
		end
		if R > 0
			d = Binomial(R,sigmaD*R/N)
			dD[i] = rand(d)
		else
			dD[i] =0
		end
		C +=  dC[i] - dR[i] - dD[i]
		R += dR[i] - dD[i]
		I += dIp - dIm - dC[i]
		S -=  dIp
	end

	return dC, dR, dD

end


### A Pluto.jl notebook ###
# v0.14.3

using Markdown
using InteractiveUtils

# ╔═╡ 4269cea4-79ac-11eb-20fb-6d35207dbbf8
begin
	using Revise
	using LinearAlgebra
	using BSplineKit
	using SparseArrays
	using Plots
end

# ╔═╡ cb24bdf0-5ba4-4683-8fcc-a635192bc173
using BenchmarkTools

# ╔═╡ 5a37f162-79ae-11eb-18fe-b329714a6617
f(x) = exp(-x / 2π) * cos(4x)

# ╔═╡ 6899fcb2-79c9-11eb-0d95-cb916a6751d3
xlims = 0, 2π

# ╔═╡ 4476d184-79ae-11eb-135f-af80be0c9bde
xdata = range(xlims...; length = 512 + 1)

# ╔═╡ 911d13ee-79b1-11eb-33b4-5d7885069097
xbreaks = range(xlims...; length = 32 + 1)

# ╔═╡ 81a5efec-79ae-11eb-01f1-31bc96f18a6d
ydata = map(x -> f(x) + 0.5 * randn(), xdata)

# ╔═╡ 22639830-79af-11eb-3834-fd118a685500
B0 = BSplineBasis(BSplineOrder(4), xbreaks)

# ╔═╡ 823cb81a-79c6-11eb-1e4d-75eb02469abf
# Impose u' = 0 at the boundaries
B = RecombinedBSplineBasis(Derivative(1), B0)
# B = B0

# ╔═╡ 2b263ef0-79af-11eb-3243-eb0ac08d1a46
C = collocation_matrix(B, xdata, SparseMatrixCSC{Float64});

# ╔═╡ 98f5a344-79af-11eb-1049-05b3511b5760
F = qr(C);

# ╔═╡ 6046792e-79af-11eb-1338-01159eefb4f3
coefs = F \ ydata

# ╔═╡ 4ed10b90-9f21-4f86-9e02-b0489c879dd1
coefs_alt = cholesky(C' * C) \ (C' * ydata)

# ╔═╡ 921213fd-fdbc-4f4b-b379-1b0432ce056e
norm(coefs - coefs_alt, Inf)

# ╔═╡ 6b405606-79af-11eb-3d80-8bdae47a03ea
fspline = Spline(B, coefs)

# ╔═╡ b76dc390-26fc-4f4e-837d-1758ea067d89
md"# Benchmarks"

# ╔═╡ 9f86c2f9-e5c3-4c9b-9c23-a93b708ce6ad
begin
	lsq_v1(C, ydata) = qr(C) \ ydata
	lsq_v2(C, ydata) = cholesky(C' * C) \ (C' * ydata)
end

# ╔═╡ 546eab20-b938-40f5-81a1-344e5d2633c9
@benchmark lsq_v1($C, $ydata)

# ╔═╡ 379d0576-0969-4e8a-b743-3ee4408d2ec7
@benchmark lsq_v2($C, $ydata)

# ╔═╡ 6c65dff6-79c3-11eb-15ab-c9757e03c434
md"""## Least-squares as in de Boor

Similar to de Boor's chapter on least-squares approximation, with the difference that we use the continuous inner product (Galerkin matrix) on the left-hand side.

"""

# ╔═╡ 999dc716-79c5-11eb-3c07-93e12d6f0241
M = galerkin_matrix(B);

# ╔═╡ 3e88842a-79c4-11eb-12a7-9fb8f940971b
function lsq_weighted_data(xs, ys)
	gs = similar(ys)
	gs[begin] = ys[begin] * (xs[begin + 1] - xs[begin]) / 2
	is = eachindex(xs)[2:end-1]
	for i in is
		gs[i] = ys[i] * (xs[i + 1] - xs[i - 1]) / 2
	end
	gs[end] = ys[end] * (xs[end] - xs[end - 1]) / 2
	gs
end

# ╔═╡ 7be65118-79c3-11eb-09de-2f1fc37ef07d
gs = lsq_weighted_data(xdata, ydata)

# ╔═╡ 7f765e48-79c5-11eb-0730-43153731b317
rhs = C' * gs

# ╔═╡ a85cda4e-79c5-11eb-2a78-f1de71c9bc30
lsq_coefs = cholesky(M) \ rhs

# ╔═╡ b97a005e-79c5-11eb-1f41-3f5a4bd0d263
f_lsq = Spline(B, lsq_coefs)

# ╔═╡ 95461eda-79ae-11eb-3d0c-5bdb5411c9c9
begin
	scatter(xdata, ydata; label = "data", markersize=3, alpha = 0.2)
	plot!(f, xlims...; colour = :black, label = "f(x)", lw = 3)
	plot!(x -> fspline(x), xlims...; label = "Fit", lw = 4)
	plot!(x -> f_lsq(x), xlims...; ls = :dash, label = "Fit LSQ", lw = 4)
end

# ╔═╡ 012e964a-79ca-11eb-3eb6-5356e32db238
begin
	plot(x -> f_lsq(x) - f(x), xlims...)
	plot!(x -> fspline(x) - f(x), xlims...)
	scatter!(xbreaks, zeros(length(xbreaks)))
end

# ╔═╡ 1c0c5472-79ca-11eb-28ee-196377c2d294
begin
	plot(x -> (f_lsq(x) - fspline(x))^2, xlims...; yscale = :log10)
	scatter!(xbreaks, fill(1e-10, length(xbreaks)))
end

# ╔═╡ Cell order:
# ╠═5a37f162-79ae-11eb-18fe-b329714a6617
# ╠═6899fcb2-79c9-11eb-0d95-cb916a6751d3
# ╠═4476d184-79ae-11eb-135f-af80be0c9bde
# ╠═911d13ee-79b1-11eb-33b4-5d7885069097
# ╠═81a5efec-79ae-11eb-01f1-31bc96f18a6d
# ╠═22639830-79af-11eb-3834-fd118a685500
# ╠═823cb81a-79c6-11eb-1e4d-75eb02469abf
# ╠═2b263ef0-79af-11eb-3243-eb0ac08d1a46
# ╠═98f5a344-79af-11eb-1049-05b3511b5760
# ╠═6046792e-79af-11eb-1338-01159eefb4f3
# ╠═4ed10b90-9f21-4f86-9e02-b0489c879dd1
# ╠═921213fd-fdbc-4f4b-b379-1b0432ce056e
# ╠═6b405606-79af-11eb-3d80-8bdae47a03ea
# ╠═95461eda-79ae-11eb-3d0c-5bdb5411c9c9
# ╟─b76dc390-26fc-4f4e-837d-1758ea067d89
# ╠═cb24bdf0-5ba4-4683-8fcc-a635192bc173
# ╠═9f86c2f9-e5c3-4c9b-9c23-a93b708ce6ad
# ╠═546eab20-b938-40f5-81a1-344e5d2633c9
# ╠═379d0576-0969-4e8a-b743-3ee4408d2ec7
# ╟─6c65dff6-79c3-11eb-15ab-c9757e03c434
# ╠═7be65118-79c3-11eb-09de-2f1fc37ef07d
# ╠═7f765e48-79c5-11eb-0730-43153731b317
# ╠═999dc716-79c5-11eb-3c07-93e12d6f0241
# ╠═a85cda4e-79c5-11eb-2a78-f1de71c9bc30
# ╠═b97a005e-79c5-11eb-1f41-3f5a4bd0d263
# ╠═012e964a-79ca-11eb-3eb6-5356e32db238
# ╠═1c0c5472-79ca-11eb-28ee-196377c2d294
# ╠═3e88842a-79c4-11eb-12a7-9fb8f940971b
# ╠═4269cea4-79ac-11eb-20fb-6d35207dbbf8


#finitemps.jl
struct FiniteMPS{S<:EuclideanSpace} <: AbstractMPS{S}
    Ts::Vector{<:AbstractTensorMap{S,2,1}}
    Ss::Vector{<:AbstractTensorMap{S,1,1}}
    canonical_form::Vector{Symbol}
    function FiniteMPS{S}(Ts::Vector{<:AbstractTensorMap{S,2,1}},
                            Ss::Vector{<:AbstractTensorMap{S,1,1}},
                            canonical_form::Vector{Symbol}) where {S<:EuclideanSpace}
        (dim(space(Ts[1],1)) == 1 && dim(space(Ts[end],3)) == 1) ||
            error("The left bond of the first tensor and the right bonnd of the
                    last tensor should have dimension 1 for finite MPS!")
        canonical_form[1] ∈ [:left_canonical,:right_canonical,:site_canonical,
                                :bond_canonical,:GammaLambda_canonical,:non_canonical] ||
            error("There is no such canonical form!")
        for i in 1:length(Ts)-1
            space(Ts[i],3) == dual(space(Ts[i+1],1)) ||
                error("The contracted spaces of tensors do not match!")
        end
        new(Ts,Ss,canonical_form)
    end
end

FiniteMPS(Ts,Ss,canonical_form) = FiniteMPS{spacetype(Ts[1])}(Ts, Ss, canonical_form)

function make_left_canonical(mps::FiniteMPS)
    for i in 1:length(mps.Ts)-1
        Q, R = leftorth(mps.Ts[i])
        mps.Ts[i] = Q
        A_tmp = TensorMap(undef, codomain(R)*codomain(mps.Ts[i+1],2), domain(mps.Ts[i+1]))
        @tensor A_tmp[a b;c] = R[a d]*mps.Ts[i+1][d b;c]
        mps.Ts[i+1] = A_tmp
    end
    mps_norm = tr(mps.Ts[L]'*mps.Ts[L])
    mps.Ts[L] = mps.Ts[L]/sqrt(mps_norm)
    mps.canonical_form[1] = :left_canonical
    return
end

function make_right_canonical(mps::FiniteMPS)
    for i in length(mps.Ts):-1:2
        L, Q = rightorth(mps.Ts[i], (1,), (2,3))
        mps.Ts[i] = permute(Q,(1,2),(3,))
        A_tmp = TensorMap(undef, codomain(mps.Ts[i-1]), domain(L))
        @tensor A_tmp[a b;d] = mps.Ts[i-1][a b;c]*L[c;d]
        mps.Ts[i-1] = A_tmp
    end
    T_tmp = TensorMap(undef, codomain(mps.Ts[1],1),codomain(mps.Ts[1],1))
    @tensor T_tmp[a;d] := mps.Ts[1][a b;c]*conj(mps.Ts[1][d b;c])
    mps_norm = tr(T_tmp)
    mps.Ts[1] = mps.Ts[1]/sqrt(mps_norm)
    mps.canonical_form[1] = :right_canonical
    return
end

function LinearAlgebra.norm(mps::FiniteMPS)
    A = mps.Ts[1]'*mps.Ts[1]
    for i in 2:length(mps.Ts)
        A_tmp1 = TensorMap(undef,codomain(A)*codomain(mps.Ts[i],2),domain(mps.Ts[i]))
        @tensor A_tmp1[d b;c] = A[d;a]*mps.Ts[i][a b;c]
        A_tmp2 = TensorMap(undef,domain(mps.Ts[i]),domain(mps.Ts[i]))
        @tensor A_tmp2[d;c] =  conj(mps.Ts[i][a b;d])*A_tmp1[a b;c]
        A = A_tmp2
    end
    return tr(A)
end

function make_site_canonical(psi::FiniteMPS, n::Int)
end

function bond_canonical(psi::FiniteMPS, b::Int)
end


"""
    GeneralizedExtremeValue(μ, σ, ξ)

The *Generalized extreme value distribution* with shape parameter `ξ`, scale `σ` and location `μ` has probability density function

```math
f(x; \\xi, \\sigma, \\mu) = \\begin{cases}
        \\frac{1}{\\sigma} \\left[ 1+\\left(\\frac{x-\\mu}{\\sigma}\\right)\\xi\\right]^{-1/\\xi-1} \\exp\\left\\{-\\left[ 1+ \\left(\\frac{x-\\mu}{\\sigma}\\right)\\xi\\right]^{-1/\\xi} \\right\\} & \\text{for } \\xi \\neq 0  \\\\
        \\frac{1}{\\sigma} \\exp\\left\\{-\\frac{x-\\mu}{\\sigma}\\right\\} \\exp\\left\\{-\\exp\\left[-\\frac{x-\\mu}{\\sigma}\\right]\\right\\} & \\text{for } \\xi = 0 \\\\
    \\end{cases}
```

for

```math
x \\in \\begin{cases}
        \\left[ \\mu - \\frac{\\sigma}{\\xi}, + \\infty \\right) & \\text{for } \\xi > 0 \\\\
        \\left( - \\infty, + \\infty \\right) & \\text{for } \\xi = 0 \\\\
        \\left( - \\infty, \\mu - \\frac{\\sigma}{\\xi} \\right] & \\text{for } \\xi < 0
    \\end{cases}
```

```julia
GeneralizedExtremeValue(μ, σ, ξ)      # Generalized Pareto distribution with shape ξ, scale σ and location μ.

params(d)       # Get the parameters, i.e. (μ, σ, ξ)
location(d)     # Get the location parameter, i.e. μ
scale(d)        # Get the scale parameter, i.e. σ
shape(d)        # Get the shape parameter, i.e. ξ (sometimes called c)
```

External links

* [Generalized extreme value distribution on Wikipedia](https://en.wikipedia.org/wiki/Generalized_extreme_value_distribution)

"""
struct GeneralizedExtremeValue{T<:Real} <: ContinuousUnivariateDistribution
    μ::T
    σ::T
    ξ::T

    function GeneralizedExtremeValue{T}(μ::T, σ::T, ξ::T) where T
        σ > zero(σ) || error("Scale must be positive")
        new{T}(μ, σ, ξ)
    end
end

GeneralizedExtremeValue(μ::T, σ::T, ξ::T) where {T<:Real} = GeneralizedExtremeValue{T}(μ, σ, ξ)
GeneralizedExtremeValue(μ::Real, σ::Real, ξ::Real) = GeneralizedExtremeValue(promote(μ, σ, ξ)...)
function GeneralizedExtremeValue(μ::Integer, σ::Integer, ξ::Integer)
    return GeneralizedExtremeValue(float(μ), float(σ), float(ξ))
end

#### Conversions
function convert(::Type{GeneralizedExtremeValue{T}}, μ::Real, σ::Real, ξ::Real) where T<:Real
    GeneralizedExtremeValue(T(μ), T(σ), T(ξ))
end
function convert(::Type{GeneralizedExtremeValue{T}}, d::GeneralizedExtremeValue{S}) where {T <: Real, S <: Real}
    GeneralizedExtremeValue(T(d.μ), T(d.σ), T(d.ξ))
end

minimum(d::GeneralizedExtremeValue{T}) where {T<:Real} =
        d.ξ > 0 ? d.μ - d.σ / d.ξ : -T(Inf)
maximum(d::GeneralizedExtremeValue{T}) where {T<:Real} =
        d.ξ < 0 ? d.μ - d.σ / d.ξ : T(Inf)

#### Parameters

shape(d::GeneralizedExtremeValue) = d.ξ
scale(d::GeneralizedExtremeValue) = d.σ
location(d::GeneralizedExtremeValue) = d.μ
params(d::GeneralizedExtremeValue) = (d.μ, d.σ, d.ξ)
@inline partype(d::GeneralizedExtremeValue{T}) where {T<:Real} = T

#### Statistics

testfd(d::GeneralizedExtremeValue) = d.ξ^3
g(d::GeneralizedExtremeValue, k::Real) = gamma(1 - k * d.ξ) # This should not be exported.

function median(d::GeneralizedExtremeValue)
    (μ, σ, ξ) = params(d)

    if abs(ξ) < eps() # ξ == 0
        return μ - σ * log(log(2))
    else
        return μ + σ * (log(2) ^ (- ξ) - 1) / ξ
    end
end

function mean(d::GeneralizedExtremeValue{T}) where T<:Real
    (μ, σ, ξ) = params(d)

    if abs(ξ) < eps(one(ξ)) # ξ == 0
        return μ + σ * MathConstants.γ
    elseif ξ < 1
        return μ + σ * (gamma(1 - ξ) - 1) / ξ
    else
        return T(Inf)
    end
end

function mode(d::GeneralizedExtremeValue)
    (μ, σ, ξ) = params(d)

    if abs(ξ) < eps(one(ξ)) # ξ == 0
        return μ
    else
        return μ + σ * ((1 + ξ) ^ (-ξ) - 1) / ξ
    end
end

function var(d::GeneralizedExtremeValue{T}) where T<:Real
    (μ, σ, ξ) = params(d)

    if abs(ξ) < eps(one(ξ)) # ξ == 0
        return σ ^2 * π^2 / 6
    elseif ξ < 1/2
        return σ^2 * (g(d, 2) - g(d, 1) ^ 2) / ξ^2
    else
        return T(Inf)
    end
end

function skewness(d::GeneralizedExtremeValue{T}) where T<:Real
    (μ, σ, ξ) = params(d)

    if abs(ξ) < eps(one(ξ)) # ξ == 0
        return 12sqrt(6) * zeta(3) / pi ^ 3 * one(T)
    elseif ξ < 1/3
        g1 = g(d, 1)
        g2 = g(d, 2)
        g3 = g(d, 3)
        return sign(ξ) * (g3 - 3g1 * g2 + 2g1^3) / (g2 - g1^2) ^ (3/2)
    else
        return T(Inf)
    end
end

function kurtosis(d::GeneralizedExtremeValue{T}) where T<:Real
    (μ, σ, ξ) = params(d)

    if abs(ξ) < eps(one(ξ)) # ξ == 0
        return T(12)/5
    elseif ξ < 1 / 4
        g1 = g(d, 1)
        g2 = g(d, 2)
        g3 = g(d, 3)
        g4 = g(d, 4)
        return (g4 - 4g1 * g3 + 6g2 * g1^2 - 3 * g1^4) / (g2 - g1^2)^2 - 3
    else
        return T(Inf)
    end
end

function entropy(d::GeneralizedExtremeValue)
    (μ, σ, ξ) = params(d)
    return log(σ) + MathConstants.γ * ξ + (1 + MathConstants.γ)
end

function quantile(d::GeneralizedExtremeValue, p::Real)
    (μ, σ, ξ) = params(d)

    if abs(ξ) < eps(one(ξ)) # ξ == 0
        return μ + σ * (-log(-log(p)))
    else
        return μ + σ * ((-log(p))^(-ξ) - 1) / ξ
    end
end

#### Support

insupport(d::GeneralizedExtremeValue, x::Real) = minimum(d) <= x <= maximum(d)

#### Evaluation

function logpdf(d::GeneralizedExtremeValue{T}, x::Real) where T<:Real
    if x == -Inf || x == Inf || ! insupport(d, x)
      return -T(Inf)
    else
        (μ, σ, ξ) = params(d)

        z = (x - μ) / σ # Normalise x.
        if abs(ξ) < eps(one(ξ)) # ξ == 0
            t = z
            return -log(σ) - t - exp(-t)
        else
            if z * ξ == -1 # Otherwise, would compute zero to the power something.
                return -T(Inf)
            else
                t = (1 + z * ξ) ^ (-1/ξ)
                return - log(σ) + (ξ + 1) * log(t) - t
            end
        end
    end
end

function pdf(d::GeneralizedExtremeValue{T}, x::Real) where T<:Real
    if x == -Inf || x == Inf || ! insupport(d, x)
        return zero(T)
    else
        (μ, σ, ξ) = params(d)

        z = (x - μ) / σ # Normalise x.
        if abs(ξ) < eps(one(ξ)) # ξ == 0
            t = exp(-z)
            return (t * exp(-t)) / σ
        else
            if z * ξ == -1 # In this case: zero to the power something.
                return zero(T)
            else
                t = (1 + z*ξ)^(- 1 / ξ)
                return (t^(ξ + 1) * exp(- t)) / σ
            end
        end
    end
end

function logcdf(d::GeneralizedExtremeValue, x::Real)
    μ, σ, ξ = params(d)
    z = (x - μ) / σ
    return if abs(ξ) < eps(one(ξ)) # ξ == 0
        -exp(- z)
    else
        # y(x) = y(bound) = 0 if x is not in the support with lower/upper bound
        y = max(1 + z * ξ, 0)
        - y^(-1/ξ)
    end
end
cdf(d::GeneralizedExtremeValue, x::Real) = exp(logcdf(d, x))

ccdf(d::GeneralizedExtremeValue, x::Real) = - expm1(logcdf(d, x))
logccdf(d::GeneralizedExtremeValue, x::Real) = log1mexp(logcdf(d, x))

#### Sampling
function rand(rng::AbstractRNG, d::GeneralizedExtremeValue)
    (μ, σ, ξ) = params(d)

    # Generate a Float64 random number uniformly in (0,1].
    u = 1 - rand(rng)

    if abs(ξ) < eps(one(ξ)) # ξ == 0
        rd = - log(- log(u))
    else
        rd = expm1(- ξ * log(- log(u))) / ξ
    end

    return μ + σ * rd
end


@afgc function _stackDftReg(resource::ArrayFireLibs,imgRef::AFArray{Complex{T},N1},imgF::AFArray{Complex{T},N},usfac) where {T,N,N1}
    if N1 != (N - 1)
        error("Reference image has the wrong dimensionality")
    end
    nimg = size(imgF,N)
    if usfac==0
        ## Compute error for no pixel shift
        @afgc CCmax = reshape(sum(imgRef.*conj(imgF),1:N1),nimg)
        @afgc rfzero = sum(abs2(imgRef))
        @afgc rgzero = reshape(sum(abs2(imgF),1:N1),nimg)
        @afgc error = 1 - CCmax.*conj(CCmax)./(rgzero*rfzero)
        @afgc diffphase = atan2(imag(CCmax),real(CCmax),false)
        output = error
    elseif usfac==1
        ## Whole-pixel shift - Compute crosscorrelation by an IFFT and locate the peak
        L = length(imgRef)
        if N1 == 2
            @afgc CC = ifft2(imgRef.*conj(imgF))
        elseif N1 == 3
            @afgc CC = ifft3(imgRef.*conj(imgF))
        end

        @afgc CC =  abs(CC)
        CCmax = maximum(CC,1)

        for i in 2:N1
            CCmax = maximum(CCmax,i)
        end

        loc = find(CC .== CCmax) ## Problem if duplicates in a frame

        @afgc rfzero = sum(abs2(imgRef))/L
        @afgc rgzero = reshape(sum(abs2(imgF),1:N1),nimg)/L
        @afgc error = abs(1 - CCmax.*conj(CCmax)./(rgzero*rfzero))
        @afgc diffphase = Array(vec(atan2(imag(CCmax),real(CCmax),false)))

        indi = size(imgRef)
        ind2 = tuple([div(x,2) for x in indi]...)

        locTemp = ind2sub(size(CC),Int.(Array(loc)))

        locI = [tuple([locTemp[j][i] for j in 1:N1]...) for i in 1:nimg]

        shift = zeros(nimg,length(locI[1]))

        for im in eachindex(locI)
            for i in eachindex(locI[im])
                if locI[im][i]>ind2[i]
                    shift[im,i]=locI[im][i]-indi[i]-1
                else shift[im,i]=locI[im][i]-1
                end
            end
        end

        #output = Dict("error"=>error,"shift"=>shift,"diffphase"=>diffphase)
        output= (shift,diffphase,error)
    else

        @afgc CC = imgRef.*conj(imgF)

        @afgc CC = padAFArray(CC,N1)

        #CC[ranges...] = imgRef.*conj(imgF)
        #afgc()
        ##  Compute crosscorrelation and locate the peak
        if N1 == 2
            @afgc CC = ifft2(CC)
        elseif N1 == 3
            @afgc CC = ifft3(CC)
        end

        @afgc CC =  abs(CC)
        CCmax = maximum(CC,1)

        for i in 2:N1
           @afgc CCmax = maximum(CCmax,i)
        end

        loc = find(CC .== CCmax)
        indi = size(CC)[1:N1]
        ind2 = tuple([div(x,2) for x in indi]...)
        locTemp = ind2sub(size(CC),Int.(Array(loc)))

        locI = [tuple([locTemp[j][i] for j in 1:N1]...) for i in 1:nimg]

        shift = zeros(nimg,length(locI[1]))

        for im in eachindex(locI)
            for i in eachindex(locI[im])
                if locI[im][i]>ind2[i]
                    shift[im,i]=locI[im][i]-indi[i]-1
                else shift[im,i]=locI[im][i]-1
                end
            end
        end

        shift = shift/2

        ## If upsampling > 2, then refine estimate with matrix multiply DFT
        if usfac > 2
            ### DFT Computation ###
            # Initial shift estimate in upsampled grid
            shift = round.(Integer,shift*usfac)/usfac
            dftShift = div(ceil(usfac*1.5),2) ## center of output array at dftshift+1
            ## Matrix multiplies DFT around the current shift estimate
            @afgc CC = conj(dftups(imgF.*conj(imgRef),ceil(Integer,usfac*1.5),usfac,dftShift-shift*usfac,N1))/(prod(ind2)*usfac^N1)
            ## Locate maximum and map back to original pixel grid

            @afgc CC =  abs(CC)
            CCmax = maximum(CC,1)

            for i in 2:N1
                @afgc CCmax = maximum(CCmax,i)
            end

            loc = find(CC .== CCmax)
            indi = size(CC)[1:N1]
            locTemp = ind2sub(size(CC),Int.(Array(loc)))

            locI = [tuple([locTemp[j][i] for j in 1:N1]...) for i in 1:nimg]

            rg00 = Array(dftups(imgRef.*conj(imgRef),1,usfac,zeros(1,N1),N1))[1]/(prod(ind2)*usfac^N)
            rf00 = Array(dftups(imgF.*conj(imgF),1,usfac,zeros(nimg,N1),N1))[repeat([1],inner=[N1])...,:]./(prod(ind2)*usfac^N)

            locI = map((x) -> x .- dftShift .- 1,locI)
            for im in eachindex(locI)
                for i in eachindex(locI[im])
                    shift[im,i]=shift[im,i]+locI[im][i]/usfac
                end
            end

        else
            rg00 = Array(sum(imgRef.*conj(imgRef)))/(prod(indi))
            rf00 = Array(sum(imgF.*conj(imgF),1:N1))/(prod(indi))
        end
        error = Array(1 - vec(CCmax).*vec(conj(CCmax)))./vec(rg00*rf00)
        error = sqrt.(abs.(error))
        diffphase = Array(vec(atan2(imag(CCmax),real(CCmax),false)))
        ## If its only one row or column the shift along that dimension has no effect. Set to zero.
        shift[:,[div(x,2) for x in size(imgRef)].==1]=0
        output = (shift,diffphase,error)
    end
output
end

@afgc function padAFArray(CC,N)
    for r in N:-1:1
        sz = collect(size(CC))
        sz[r]=div(size(CC)[r],2)
        padR = zeros(AFArray{Complex{Float64}},sz...)
        @afgc CC = cat(r,padR,CC)
        @afgc CC = cat(r,CC,padR)
    end
    CC
end

function SubpixelRegistration.stackDftReg(resource::ArrayFireLibs,imgser::AbstractArray{T,N};ref::AbstractArray{T,N1}=reshape(selectdim(imgser,N,1),size(imgser)[1:(N-1)]),ufac::Int=10,chunkSize=100) where {T,N,N1}
    if ((N1 != (N - 1)) & (N1 != N))
        error("Reference image has the wrong dimensionality")
    end
    @afgc ref = fft(AFArray(ref))
    imgF = fft(imgser,(1:N1...,))

    if N1 == (N-1)
        chunks = [(i:min(chunkSize+i-1,size(imgF)[N])) for i in 1:chunkSize:size(imgF)[N]]
        shifts = zeros(size(imgF)[N],N1)
        diffphase = zeros(size(imgF)[N])
        error = zeros(size(imgF)[N])
        for ch in chunks
            shifts[ch,:],diffphase[ch],error[ch] = _stackDftReg(ArrayFireLibs(),ref,AFArray(selectdim(imgF,N,ch)),ufac)
        end
        results = (shifts,diffphase,error)
    else
        results = _stackDftReg(ArrayFireLibs(),ref,AFArray(imgF),ufac)
    end
    results = map(1:length(diffphase)) do i
        Dict("shift" => shifts[i,:],"diffphase" => diffphase[i],"error"=>error[i])
    end
    return results
end

"
`dftups{T,N}(inp::AbstractArray{T,N},no,usfac::Int=1,off=zeros(N))`

Upsampled DFT by matrix multiplication, computes an upsampled DFT in just a small region. `no` is the size of the region in pixels, `offset` the position in the full array, `usfac` the upsampling parameter."
@afgc function dftups(inp::AFArray,no,usfac::Int=1,offset=zeros(ndims(inp)),N=ndims(inp))
    sz = [size(inp)...]
    Nf = ndims(inp)
    nim = sz[Nf]
    permV = 1:N
    permFull = 0:3
    if Nf == N
        for i in permV
            pFull = collect(permFull)
            idxs = [pFull[i];deleteat!(pFull,i)]
            @afgc inp = reorder(inp,idxs...)
            leftK = AFArray((0:(no-1)).-offset[:,i]')
            @afgc leftK = reshape(leftK,no)
            rightK = AFArray(ifftshift(0:(sz[i]-1))-floor(sz[i]/2))
            @afgc rightK = rightK'
            kern = leftK*rightK
            @afgc kern = exp.(-1im*2*pi/(sz[i]*usfac)*kern)
            nsz=size(inp)
            d = nsz[2:Nf]
            #kern = reorder(AFArray(repeat(kern,outer=(1,1,nsz[1]))),0,2,1,4)
            @afgc inp = kern * reshape(inp, Val(2))
            @afgc inp = reshape(inp,no,d...)
        end
    else
        for i in permV
            pFull = collect(permFull)
            idxs = [pFull[i];deleteat!(pFull,i)]
            @afgc inp = reorder(inp,idxs...)
            leftK = AFArray((0:(no-1)).-offset[:,i]')
            @afgc leftK = reshape(leftK,no,1,nim)
            rightK = AFArray(repeat(ifftshift(0:(sz[i]-1)),outer=(1,sz[Nf]))-floor(sz[i]/2))
            @afgc rightK = reshape(rightK,1,sz[i],sz[Nf])

            kern = leftK*rightK
            @afgc kern = exp.(-1im*2*pi/(sz[i]*usfac)*kern)

            nsz=size(inp)
            d = nsz[2:Nf]
            remDim = prod(nsz[2:N])

            @afgc inp = kern * reshape(inp, nsz[1],remDim,sz[Nf])
            @afgc inp = reshape(inp,no,d...)
        end
    end
    newIdx = collect(N:-1:1).-1
    extraIdx = setdiff(0:3,newIdx)
    reorder(inp,newIdx...,extraIdx...)
end

"
`subPixShift(imgft::AbstractArray{Complex{Float64}},shift::Array{Float64,1})`

Shift the image `imgft` (in Fourier space) by the amount provided in the vector `shift`."
@afgc function subPixShift(imgft::AFArray{Complex{T},N},shift::Array{Float64,2},diffphase::Array{Float64,1}) where {T,N}
    sz = [size(imgft)[1:(N-1)]...]
    nim = size(imgft)[N]
    Z=0
    resh = [repeat([1],inner=N-1)...,nim]
    for i in eachindex(sz)
        shifti = AFArray(ifftshift((-div(sz[i],2)):(ceil(Integer,sz[i]/2)-1)).*shift[:,i]'/sz[i])
        reshTemp = copy(resh)
        reshTemp[i] = size(shifti)[1]
        Z = Z .- reshape(shifti,reshTemp...)
    end

    Greg = imgft .* exp.(2im*pi*Z)
    @afgc Greg = Greg .* reshape(exp.(AFArray(1im*diffphase)),resh...)
    Greg
end

function subPixShift(imgft::AFArray{Complex{T}},shift::Array{Float64,1},diffphase) where {T}
    sz = [size(imgft)...]
    N=0
    for i in eachindex(sz)
        shifti = ifftshift((-div(sz[i],2)):(ceil(Integer,sz[i]/2)-1))*shift[i]/sz[i]
        resh = (repeat([1],inner=[i-1])...,length(shifti))
        N = N .- reshape(shifti,resh...)
    end

    Greg = imgft .* exp.(2im*pi*N)
    Greg = Greg .* exp(1im*diffphase)
    Greg
end

function SubpixelRegistration.alignFromDict(resource::ArrayFireLibs,img2reg::AbstractArray{T,N},dftRegRes::Array{Dict{String,Any},1}) where {T,N}
    if (length(dftRegRes) != size(img2reg)[N])
        error("Alignment results and image stack dimensionalities don't match.")
    end

    img2regF = fft(img2reg,(1:(N-1)...,))
    szF = size(img2reg)[1:(N-1)]

    imRes = subPixShift(AFArray(img2regF),SubpixelRegistration.shifts_from_res(dftRegRes),SubpixelRegistration.diffphase_from_res(dftRegRes))

    imRes = real(ifft(Array(imRes),1:(N-1)))

    imRes
end


# Adapted from:
# SOSDEMO4 --- Matrix Copositivity
# Section 3.4 of SOSTOOLS User's Manual

@testset "SOSDEMO4 with $(factory.constructor)" for factory in sdp_factories
    @polyvar x[1:5]

    # The matrix under consideration
    J = [1 -1  1  1 -1;
        -1  1 -1  1  1;
         1 -1  1 -1  1;
         1  1 -1  1 -1;
        -1  1  1 -1  1]

    xs = vec(x).^2
    xsJxs = dot(xs, J*xs)
    r = sum(xs)

    m0 = SOSModel(factory)
    @constraint m0 xsJxs >= 0
    JuMP.optimize!(m0)
    @test JuMP.dual_status(m0) == MOI.InfeasibilityCertificate

    m1 = SOSModel(factory)
    @constraint m1 r*xsJxs >= 0
    JuMP.optimize!(m1)
    @test JuMP.primal_status(m1) == MOI.FeasiblePoint
    # Program is feasible. The matrix J is copositive
end


#
# Computes the ordered internal distance matrix of the coordinates of
# a structure
#
# Input: x[n,3] the 3D coordinates of the atoms
#
# Returns: D, the internal distance matrix of the x
#          Dind, the indices of the elements of D, ordered from lower
#          the higher at each column
#

function intDmat(x :: Matrix{Float64})

  n = length(x[:,1])
  D = Matrix{Float64}(undef,n,n)
  Dind = Matrix{Int64}(undef,n,n)

  # Compute disprot, the internal square-distance matrix

  for i in 1:n
    D[i,i] = 0.
  end

  for j in 1:n-1
    for i in j+1:n
      dij = ( x[i,1] - x[j,1] )^2 + ( x[i,2] - x[j,2] )^2 + ( x[i,3] - x[j,3] )^2
      D[i,j] = dij
      D[j,i] = dij
    end
  end

  # Sort the distances to atom i of the protein

  for j in 1:n
    ind = [ i for i in 1:n ]
    Dind[:,j] = sort(ind, by = i -> D[i,j] )
  end

  return D, Dind

end


struct Boat_ode_params
    bh2o_0::Float64
    angle_ranges::Boat_angle_ranges # in radians
    timing::Boat_timing
    x1va::Float64
    x1air::Float64
end

Boat_ode_params(θ0, bs::Boat_settings) =
    Boat_ode_params(bh2o0(θ0), deg2rad(bs.angle_ranges), bs.timing, bs.x1va, bs.x1air)

function Boat_sim_pars(z, θ, u, w, p::Boat_ode_params, t)
    sθ, cθ = sincos(θ)
    R = SA[cθ -sθ; sθ cθ]
    U = SA[u, w]
    Boat_sim_pars(
        get_bh2o(z, θ, p),
        Boat_angles(p.angle_ranges, p.timing, t),
        η(t, p.timing),
        p.x1va,
        p.x1air,
        (R * U)...,
    )
end

function boatode!(dy, y, p::Boat_ode_params, t)
    # y contains int(u,t), int(w, t), the - position level coordinates
    # as well as u, w, the1  - velocity level coordinates
    _, z, θ, u, w, θ′ = y

    bsp = Boat_sim_pars(z, θ, u, w, p, t)

    M = system_lhs(bsp.boat_angles)
    b = system_rhs(u, w, θ′, bsp.boat_angles)
    F = Q(θ, bsp)

    dy[1] = bsp.x′
    dy[2] = bsp.z′
    dy[3] = θ′
    dy[4:6] = M \ (b + F)
end

function static_equilibrium(θ, settings::Boat_settings)
    (Q_g(
        θ,
        deg2rad(settings.angle_ranges.γ_OA[1]),
        deg2rad(settings.angle_ranges.θ_k[1]),
        deg2rad(settings.angle_ranges.θ_t[1]),
    )+Q_BUOY(θ, bh2o0(θ)))[3]
end

function find_the0(themin, themax, settings::Boat_settings)
    bisectionMethod((the) -> static_equilibrium(the, settings), themin, themax, 1e-2)
end


function get_U0(γ, β, rdisp, Nₚ)
    psamples_raw = get_psamples(rdisp, Nₚ)
    s = 0.0
    for psample in psamples_raw
        εₚ⁺, εₚ⁻, wₚ = psample
        κₚ = sqrt(εₚ⁻^2 + γ^2)
        s += wₚ*(tanh(β*0.5*(εₚ⁺+κₚ)) + tanh(β*0.5*(κₚ-εₚ⁺)))/κₚ
    end
    return 4.0/s
end

check_gamma(γ::Real, params::AbstractParams, rdisp::ReducedDispersion, Nₚ) = check_gamma(γ, params, get_samples(rdisp, Nₚ))
check_gamma(γ::Real, params::AbstractParams, psamples_raw::Vector{NTuple{3, Float64}}) = check_gamma(γ, params.β, params.u, psamples_raw)

function check_gamma(γ::Real, β::Real, u::Real, psamples_raw::Vector{NTuple{3, Float64}})
    s = 0.0
    for psample in psamples_raw
        εₚ⁺, εₚ⁻, wₚ = psample
        κₚ = sqrt(εₚ⁻^2 + γ^2)
        s += wₚ*(tanh(β*0.5*(εₚ⁺+κₚ)) + tanh(β*0.5*(κₚ-εₚ⁺)))/κₚ
    end
    return 0.25*s-u
end

function get_gamma(params::AbstractParams, rdisp::ReducedDispersion, Nₚ, range)
    psamples_raw = get_psamples(rdisp, Nₚ)
    return find_zero((γ)->check_gamma(γ, params, psamples_raw), range, Bisection())
end

function get_gamma(params::AbstractParams, psamples_raw::Vector{NTuple{3, Float64}}, range)
    return find_zero((γ)->check_gamma(γ, params.β, params.u, psamples_raw), range, Bisection())
end

function get_gamma(β::Real, u::Real, psamples_raw::Vector{NTuple{3, Float64}}, range)
    return find_zero((γ)->check_gamma(γ, β, u, psamples_raw), range, Bisection())
end

calculate_sce_rhs(k, params::AbstractParams, rdisp::ReducedDispersion, Nₚ::Int64) = calculate_sce_rhs(k, params, get_psamples(rdisp, Nₚ))

function calculate_sce_rhs(k::T, params::AbstractParams, psamples_raw::AbstractVector{NTuple{3, Float64}}) where {T<:Number}
    β, m, U0 = params.β, params.m, 1/params.u
    k² = k^2
    γ = 4*K(k²)*m/β
    kk = 4k/(1+k)^2
    s = zero(T)
    @fastmath for psample in psamples_raw
        εₚ⁺, εₚ⁻, wₚ = psample
        nₚ = γ^2*k/(εₚ⁻^2+0.25*γ^2*(1+k)^2)
		κₚ = sqrt((εₚ⁻^2+0.25*γ^2*(1-k)^2)/(1.0+0.25*(γ/εₚ⁻)^2*(1+k)^2))*Pi(nₚ, pi/2, kk)/K(kk)
        s += wₚ*(tanh(β*0.5*(εₚ⁺+κₚ)) + tanh(β*0.5*(κₚ-εₚ⁺)))*abs(εₚ⁻)/sqrt((εₚ⁻^2+0.25*γ^2*(1+k)^2)*(εₚ⁻^2+0.25*γ^2*(1-k)^2))
    end
    return 1.0 - 0.25*U0*s
end

function solve_sce_nob1(params::AbstractParams, psamples_raw::Vector{NTuple{3, Float64}}, range = (0.5, 0.99999))
    k = find_zero((x)->calculate_sce_rhs(x, params, psamples_raw), range, Bisection())
    γ = 4*K(k^2)*params.m/params.β
    return k, γ
end

function _seed_sn(params::AbstractParams, psamples_raw::Vector{NTuple{3, Float64}}, range = (0.5, 0.99999))
    β, m, N, N′ = params.β, params.m, params.N, get_length(params)
    k, γ = solve_sce_nob1(params, psamples_raw, range)
    dτ = β/(N*m)
    τs = [dτ*(i-1) for i = 1:N]
    sns = zeros(N′)
    sns[1:N] .= k*γ*Jacobi.sn.(γ*τs, k^2)
    pin_bs!(sns, params)
    return sns
end

function _seed_sn_asymp(params::AbstractParams, psamples_raw::Vector{NTuple{3, Float64}}, range = (0.001, 1000.0))
    β, m, N, N′ = params.β, params.m, params.N, get_length(params)
    γ = get_gamma(params, psamples_raw, range)
    dτ = β/(N*m)
    τs = [dτ*(i-1) for i = 1:N]
    sns = zeros(N′)
    N1 = div(N, 4)
    N2 = div(N, 2)
    sns[1:N1] .= γ.*tanh.(γ.*τs[1:N1])
    sns[(N1+1):N2] .= γ.*tanh.(γ.*(τs[N2+1].-τs[(N1+1):N2]))
    sns[(N2+1):N] .= -sns[1:N2]
    pin_bs!(sns, params)
    return sns
end

function get_static_fmin(β::Real, u::Real, psamples_raw::Vector{NTuple{3, Float64}}, range = (0.001, 1000.0))
	γ = get_gamma(β, u, psamples_raw, range)
	γ′ = Float128(γ)
	β′ = Float128(β)
	s = β*u*γ^2
	for (εₚ⁺, εₚ⁻, wₚ) in psamples_raw
		εₚ⁻′ = Float128(εₚ⁻)
		εₚ⁺′ = Float128(εₚ⁺)
		κₚ′ = sqrt(γ′^2 + εₚ⁻′^2)
		s -= wₚ*Float64(log((cosh(β′*κₚ′) + cosh(β′*εₚ⁺′))/(cosh(β′*εₚ⁻′) + cosh(β′*εₚ⁺′))))
	end
	return γ, s
end

seed_sn(params::AbstractParams, rdisp::ReducedDispersion, Nₚ::Int64) = seed_sn(params, get_psamples(rdisp, Nₚ))
function seed_sn(params::AbstractParams, psamples_raw::Vector{NTuple{3, Float64}})
    try
        _seed_sn(params, psamples_raw)
    catch e
        _seed_sn_asymp(params, psamples_raw)
    end
end

function get_τs(params::AbstractParams)
    β, m, N = params.β, params.m, params.N
    dτ = β/(N*m)
    return [dτ*(i-1) for i = 1:N]
end
function get_τs(β::Real, m::Int64, N::Int64)
    dτ = β/(N*m)
    return [dτ*(i-1) for i = 1:N]
end

show_sn(params::AbstractParams, rdisp::ReducedDispersion, Nₚ, range = (0.5, 0.99999)) = get_τs(params), seed_sn(params, rdisp, Nₚ, range)

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

seed_const(c, params::TTCrystal.AbstractParamsNoB1) = fill!(Vector{Float64}(undef, params.N),c)
function seed_const(c1, c2, params::AbstractParamsB1)
    N = params.N
    bs = Vector{Float64}(undef, 2N)
    bs[1:N] .= c1
    bs[(N+1):2N] .= c2
    return bs
end

function seed_const(c1, c2, c3, params::Union{ParamsB1B3, ParamsB3})
    N = params.N
	bs = Vector{Float64}(undef, get_length(params))
    bs[1:N] .= c1
    bs[(N+1):2N] .= c2
	bs[(2N+1):3N] .= c3
    return bs
end

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

isB1(params::AbstractParamsB1) = Val(true)
isB1(params::AbstractParamsNoB1) = Val(false)

function show_bs(res::Optim.MultivariateOptimizationResults, ::Val{true})
    bs = Optim.minimizer(res)
    N = div(length(bs), 2)
    b0 = bs[1:N]
    b1 = bs[(N+1):2N]
    plot(b0)
    plot!(b1)
end
function show_bs(res::Optim.MultivariateOptimizationResults, ::Val{false})
    bs = Optim.minimizer(res)
    plot(bs)
end

show_bs(res, params::AbstractParams) = show_bs(res, isB1(params))

function check_derivative(bs)
    N = div(length(bs), 2)
    b0 = bs[1:N]
    b1 = bs[(N+1):2N]
    b_der = similar(b0)
    b_der[1] = b0[2]-b0[end]
    b_der[end] = b0[1]-b0[end-1]
    for i=2:(N-1)
        b_der[i] = b0[i+1] - b0[i-1]
    end
    #b_der .*= -sum(abs2, b1)/sum(abs2, b_der)
	b_der .*= b1[div(N,2)+1]/b_der[div(N,2)+1]
    plot(b_der)
    plot!(b1)
end
check_derivative(res::Optim.MultivariateOptimizationResults) = check_derivative(Optim.minimizer(res))

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

function el_part(γ, P, β, psamples_raw)
	β′ = Float128(β)
    γ′ = Float128(γ)
    s = 0.0
	for (εₚ⁺, εₚ⁻, wₚ) in psamples_raw
		εₚ⁻′ = Float128(εₚ⁻)
		εₚ⁺′ = Float128(εₚ⁺)
        κₚ′ = sqrt(γ′^2 + (εₚ⁻′ - P)^2)
		s += wₚ*Float64(log((cosh(β′*κₚ′) + cosh(β′*εₚ⁺′))/(cosh(β′*εₚ⁻′) + cosh(β′*εₚ⁺′))))
	end
    return s
end

function static_f_energy(γ, P, β, u, u₃, psamples_raw)
    β*(u*γ^2 + u₃*P^2) - el_part(γ, P, β, psamples_raw)
end

function energy_grad!(g, v, β, u, u₃, psamples_raw)
    γ, P = v[1], v[2]
    g[1] = 2*β*u*γ
    g[2] = 2*β*u₃*P
    s = 0.0
    s1 = 0.0
	for (εₚ⁺, εₚ⁻, wₚ) in psamples_raw
        κₚ = sqrt(γ^2 + (εₚ⁻-P)^2)
        s += wₚ*(tanh(0.5*β*(κₚ+εₚ⁺)) + tanh(0.5*β*(κₚ-εₚ⁺)))/κₚ
        s1 += wₚ*(εₚ⁻-P)*(tanh(0.5*β*(κₚ+εₚ⁺)) + tanh(0.5*β*(κₚ-εₚ⁺)))/κₚ
    end
    g[1] -= β*0.5*γ*s
    g[2] += β*0.5*s1
end

function optimize_γ_P(γ₀, P₀, β, u, u₃, psamples_raw)
    f(v) = static_f_energy(v[1], v[2], β, u, u₃, psamples_raw)
    g!(g, v) = energy_grad!(g, v, β, u, u₃, psamples_raw)
    x0 = [γ₀, P₀]
    return optimize(f, g!, x0, Optim.Options(g_tol = 1e-10, show_trace = true))
end


# Samuel J Palmer
# Based on "A New Method of Interpolation and Smooth Curve Fitting Based on Local Parameters", Akima, 1970

module CubicSplines

function extrapolate(xs, ys)
    # Quadratic extrapolation of (x4,y4), (x5,y5)
    x1, x2, x3 = xs
    y1, y2, y3 = ys
    x4 = x3-x1+x2
    x5 = x3-x1+x3
    y4 = y3 + 2*(x4-x3)*(y3-y2)/(x3-x2) - (x4-x3)*(y2-y1)/(x2-x1)
    y5 = y4 + 2*(x5-x4)*(y4-y3)/(x4-x3) - (x5-x4)*(y3-y2)/(x3-x2)
    return x4, x5, y4, y5
end

struct CubicSpline
    xs
    ys
    ts
    logx
    function CubicSpline(xs, ys; logx=false)
        @assert length(xs) == length(ys)
        @assert issorted(xs)
        if !issorted(xs)
            idx = sortperm(xs)
            xs = xs[idx]
            ys = ys[idx]
        end
        if logx
            xs = log.(xs)
        end
        xb, xa, yb, ya = extrapolate(xs[3:-1:1], ys[3:-1:1])
        xy, xz, yy, yz = extrapolate(xs[end-2:end], ys[end-2:end])
        xs = [xa; xb; xs; xy; xz]
        ys = [ya; yb; ys; yy; yz]
        ms = zeros(eltype(ys), length(xs)-1)
        for i in 1:length(ms)-1
            ms[i] = (ys[i+1]-ys[i]) ./ (xs[i+1]-xs[i])
        end
        ts = zeros(eltype(ys), length(xs))
        for i in 3:length(xs)-2
            m1, m2, m3, m4 = ms[i-2:i+1]
            if m1 == m2 && m3 == m4
                ts[i] = (m2+m3)/2
            else
                numer = abs(m4-m3)*m2 + abs(m2-m1)*m3
                denom = abs(m4-m3)    + abs(m2-m1)
                ts[i] = numer ./ denom
            end
        end
        return new(xs[3:end-2], ys[3:end-2], ts[3:end-2], logx)
    end
end

function Base.getindex(spline::CubicSpline, x::Real)
    if spline.logx
        x = log(x)
    end
    if x < spline.xs[1]
        throw("x too small ($x < $(spline.xs[1]))")
    elseif x > spline.xs[end]
        throw("x too big ($x > $(spline.xs[end]))")
    end
    i = 1
    while !(spline.xs[i] <= x <= spline.xs[i+1])
        i += 1
    end
    x1, x2 = spline.xs[i:i+1]
    y1, y2 = spline.ys[i:i+1]
    t1, t2 = spline.ts[i:i+1]
    p0 = y1
    p1 = t1
    p2 = (3(y2-y1)/(x2-x1)-2t1-t2)/(x2-x1)
    p3 = (t1+t2-2(y2-y1)/(x2-x1))/(x2-x1)^2
    return p0 + p1*(x-x1) + p2*(x-x1)^2 + p3*(x-x1)^3
end

function Base.getindex(spline::CubicSpline, xs::AbstractArray{<:Real,1})
    return [spline[x] for x in xs]
end

function Base.firstindex(spline::CubicSpline)
    return spline.xs[1]
end

function Base.lastindex(spline::CubicSpline)
    return spline.xs[end]
end

export CubicSpline

end # module


"""
Update the parameters of the model using the gardients (∇) and the learning rate (η).
"""
function update_model_weights(parameters, ∇, η)
    L = trunc(Int, length(parameters) / 2)

    # update the parameters (weights and biases) for all the layers
    for l = 1:L
        parameters[string("W_", l)] = min.(max.(parameters[string("W_", l)] - η .* ∇[string("δW_", l)], -3), 3)
        # try not editign biases, and get swish working
        parameters[string("b_", l)] = min.(max.(parameters[string("b_", l)] - η .* ∇[string("δb_", l)], -3), 3)

        # Debugging printlns
        #println(∇[string("δW_", l)][1:5])
        #println(parameters[string("W_", l)][1:5])
    end # for

    return parameters
end # function


using OmegaCore, Distributions
using Test
import Random
Random.seed!(1)

function test_cond!()
  function f(ω)
    x = (1 ~ Normal(0, 1))(ω)
    cond!(ω, x > 0)
    x
  end
  @test all(randsample(f, 100) .>= 0.0)
end

"Test conditioning on predicate of positive measure"
function test_pos_measure()
  x = 1 ~ Normal(0, 1)
  y(ω) = x(ω) > 0
  x_cond = x |ᶜ y
  samples = randsample(x_cond, 100000; alg = RejectionSample)
  samplemean = mean(samples)
  exactmean = mean(truncated(Normal(0, 1), 0, Inf))
  @test samplemean ≈ exactmean atol = 0.01
end

"Test whether the logpdf is correct"
function test_density_cond()
  rng = Random.MersenneTwister(0)
  μ = 1 ~ Normal(0, 1)
  x = Normal(μ, 1.0)

  μ_ = -0.4321
  x_ = 0.1234
  μₓ = μ |ᶜ (x ==ₚ x_)
  ω = LazyΩ(μ => μ_)
  logpdf_ = logpdf(Normal(0, 1), μ_) + logpdf(Normal(μ_, 1), x_)
  @test logpdf(μₓ, ω) == logpdf_
end

function test_out_of_order_condition()
  x = 1 ~ Poisson(1.3)
  function f(ω)
    n = x(ω)
    x_ = 0.0
    for i = 1:n
      x_ += (i + 1 ~ Normal(0, 1))(ω)
    end
    x_
  end

  x_ = x |ᶜ (f ==ₚ 3.0)
  g_ = ω -> (x(ω), f_(ω))
end

function test_parent()
  μ = 1 ~ Normal(0, 1)
  x = 2 ~ Normal(μ, 1)
  x_ = 0.123
  μ_ = 0.1
  μ |ᶜ x == x_
  ω = defΩ()((1,) => μ_)
  @test logpdf(μ_, ω) = logpdf(Normal(0, 1), μ_) + logpdf(Normal(μ_, 1), x_)
end

function test_condition()
  μ = 1 ~ Normal(0, 1)
  x = 2 ~ Normal(μ, 1)
  x_ = 0.123
  μ_ = 0.987

  μc = μ |ᶜ x ==ₚ x_
  ω = defω()
  ω[(1,)] = μ_
  μc(ω)
  @test x(ω) == x_
  logpdf_ = logpdf(Normal(0, 1), μ) + logpdf(Normal(μ, 1), x_)
  # Will ω have a value for (2,)?
end

@testset "Conditions" begin
  test_cond!()
  test_pos_measure()
  # test_density_cond()
  # test_out_of_order_condition()
  # test_parent()
  # test_condition()
end


"""
        summary(decisionMat, weights, fns, methods)

Apply more methods for a given decision problem. The methods accept standart number of arguments.

# Arguments:
 - `decisionMat::DataFrame`: n × m matrix of objective values for n candidate (or strategy) and m criteria
 - `weights::Array{Float64, 1}`: m-vector of weights that sum up to 1.0. If the sum of weights is not 1.0, it is automatically normalized.
 - `fns::Array{Function, 1}`: m-vector of function that are either minimize or maximize.
 - `methods::Array{Symbol, 1}`: Array of symbols. The elements can be :topsis, :electre, :cocoso, :copras, :moora, :vikor, :grey, :aras, :saw, :wpm, :waspas, :edas, :marcos, :mabac, :mairca, :copras, :critic

# Description
    This method outputs a summarized output using more than MCDM methods in a comparable way.

    # Output
- `::DataFrame`: A DataFrame object, methods in columns, and alternatives in rows. Green check symbol indicates the selected alternative as the best by the corresponding method.

# Examples
```julia-repl
julia> df = DataFrame(
:age        => [6.0, 4, 12],
:size       => [140.0, 90, 140],
:price      => [150000.0, 100000, 75000],
:distance   => [950.0, 1500, 550],
:population => [1500.0, 2000, 1100]);

julia> methods = [:topsis, :electre, :vikor, :moora, :cocoso, :wpm, :waspas]

julia> w  = [0.036, 0.192, 0.326, 0.326, 0.12];

julia> fns = [maximum, minimum, maximum, maximum, maximum];

julia> result = summary(df, w, fns, methods)
3×7 DataFrame
 Row │ topsis  electre  cocoso  moora   vikor   wpm     waspas
     │ String  String   String  String  String  String  String
─────┼─────────────────────────────────────────────────────────
   1 │                           ✅      ✅
   2 │  ✅      ✅       ✅                      ✅      ✅
   3 │

```
"""
function summary(
    decisionMat::DataFrame,
    weights::Array{Float64,1},
    fns::Array{Function,1},
    methods::Array{Symbol,1})

    nmethods = length(methods)
    nalternatives, _ = size(decisionMat)

    resultdf = DataFrame()

    check = " ✅ "

    if :topsis in methods
        topresult = topsis(decisionMat, weights, fns)
        resultdf[:,:topsis] = map(x -> if topresult.bestIndex == x check else " " end, 1:nalternatives)
    end

    if :electre in methods
        # electre best index is a tuple
        result = electre(decisionMat, weights, fns)
        resultdf[:,:electre] = map(x -> if x in result.bestIndex check else " " end, 1:nalternatives)
    end

    if :cocoso in methods
        result = cocoso(decisionMat, weights, fns)
        resultdf[:,:cocoso] = map(x -> if result.bestIndex == x check else " " end, 1:nalternatives)
    end

    if :copras in methods
        result = copras(decisionMat, weights, fns)
        resultdf[:,:copras] = map(x -> if result.bestIndex == x check else " " end, 1:nalternatives)
    end

    if :moora in methods
        result = moora(decisionMat, weights, fns)
        resultdf[:,:moora] = map(x -> if result.bestIndex == x check else " " end, 1:nalternatives)
    end

    if :vikor in methods
        result = vikor(decisionMat, weights, fns)
        resultdf[:,:vikor] = map(x -> if result.bestIndex == x check else " " end, 1:nalternatives)
    end

    if :grey in methods
        result = grey(decisionMat, weights, fns)
        resultdf[:,:grey] = map(x -> if result.bestIndex == x check else " " end, 1:nalternatives)
    end

    if :aras in methods
        result = aras(decisionMat, weights, fns)
        resultdf[:,:aras] = map(x -> if result.bestIndex == x check else " " end, 1:nalternatives)
    end

    if :saw in methods
        result = saw(decisionMat, weights, fns)
        resultdf[:,:saw] = map(x -> if result.bestIndex == x check else " " end, 1:nalternatives)
    end

    if :wpm in methods
        result = wpm(decisionMat, weights, fns)
        resultdf[:,:wpm] = map(x -> if result.bestIndex == x check else " " end, 1:nalternatives)
    end

    if :waspas in methods
        result = waspas(decisionMat, weights, fns)
        resultdf[:,:waspas] = map(x -> if result.bestIndex == x check else " " end, 1:nalternatives)
    end

    if :edas in methods
        result = edas(decisionMat, weights, fns)
        resultdf[:,:edas] = map(x -> if result.bestIndex == x check else " " end, 1:nalternatives)
    end

    if :marcos in methods
        result = marcos(decisionMat, weights, fns)
        resultdf[:,:marcos] = map(x -> if result.bestIndex == x check else " " end, 1:nalternatives)
    end

    if :mabac in methods
        result = mabac(decisionMat, weights, fns)
        resultdf[:,:mabac] = map(x -> if result.bestIndex == x check else " " end, 1:nalternatives)
    end

    if :mairca in methods
        result = mairca(decisionMat, weights, fns)
        resultdf[:,:mairca] = map(x -> if result.bestIndex == x check else " " end, 1:nalternatives)
    end

    if :copras in methods
        result = copras(decisionMat, weights, fns)
        resultdf[:,:copras] = map(x -> if result.bestIndex == x check else " " end, 1:nalternatives)
    end

    if :critic in methods
        result = critic(decisionMat, fns)
        resultdf[:,:critic] = map(x -> if result.bestIndex == x check else " " end, 1:nalternatives)
    end

    return resultdf

end

"""
        summary(setting, methods)

Apply more methods for a given decision problem. The methods accept standart number of arguments.

# Arguments:
 - `setting::MCDMSetting`: MCDMSetting object.
 - `methods::Array{Symbol, 1}`: Array of symbols. The elements can be :topsis, :electre, :cocoso, :copras, :moora, :vikor, :grey, :aras, :saw, :wpm, :waspas, :edas, :marcos, :mabac, :mairca, :copras, :critic

# Description
    This method outputs a summarized output using more than MCDM methods in a comparable way.
"""
function summary(setting::MCDMSetting, methods::Array{Symbol,1})
    summary(
        setting.df,
        setting.weights,
        setting.fns,
        methods
    )
end


using OffsetArrays
using .SetProb

# ============================================================================
function rp1(parms::CParam,
             maxm::Int, meqn::Int, mwaves::Int, maux::Int, mbc::Int, mx::Int,
             ql::OffsetArray{Float64,2}, qr::OffsetArray{Float64,2},
             auxl::OffsetArray{Float64}, auxr::OffsetArray{Float64},
             wave::OffsetArray{Float64,3}, s::OffsetArray{Float64,2},
             amdq::OffsetArray{Float64,2}, apdq::OffsetArray{Float64,2})
# ============================================================================

# subroutine rp1(maxm,meqn,mwaves,maux,mbc,mx,ql,qr,auxl,auxr,wave,s,amdq,apdq)

# Riemann solver for the acoustics equations in 1d,

# On input, ql contains the state vector at the left edge of each cell
#           qr contains the state vector at the right edge of each cell

# On output, wave contains the waves,
#            s the speeds,
#
#            amdq = A^- Delta q,
#            apdq = A^+ Delta q,
#                   the decomposition of the flux difference
#                       f(qr(i-1)) - f(ql(i))
#                   into leftgoing and rightgoing parts respectively.
#

# Note that the i'th Riemann problem has left state qr(i-1,:)
#                                    and right state ql(i,:)
# From the basic clawpack routines, this routine is called with ql = qr

#    dimension wave(meqn, mwaves, 1-mbc:maxm+mbc)
#    dimension    s(mwaves,1-mbc:maxm+mbc)
#    dimension   ql(meqn, 1-mbc:maxm+mbc)
#    dimension   qr(meqn, 1-mbc:maxm+mbc)
#    dimension apdq(meqn, 1-mbc:maxm+mbc)
#    dimension amdq(meqn, 1-mbc:maxm+mbc)

#     local arrays
#     ------------
#    dimension delta(2)
delta = Vector{Float64}(undef, 2)

# density, bulk modulus, and sound speed, and impedence of medium:
# (should be set in setprob.f)
#    common /cparam/ rho,bulk,cc,zz

    # split the jump in q at each interface into waves

    # find a1 and a2, the coefficients of the 2 eigenvectors:

    for i = 2-mbc : mx+mbc
        delta[1] = ql[1,i] - qr[1,i-1]
        delta[2] = ql[2,i] - qr[2,i-1]
        a1 = (-delta[1] + parms.zz*delta[2]) / (2.0*parms.zz)
        a2 = ( delta[1] + parms.zz*delta[2]) / (2.0*parms.zz)

        # Compute the waves.

        wave[1,1,i] = -a1*parms.zz
        wave[2,1,i] = a1
        s[1,i] = -parms.cc

        wave[1,2,i] = a2*parms.zz
        wave[2,2,i] = a2
        s[2,i] = parms.cc

    end

    # compute the leftgoing and rightgoing flux differences:
    # Note s(1,i) < 0   and   s(2,i) > 0.

    for m = 1 : meqn
        for i = 2-mbc : mx+mbc
            amdq[m,i] = s[1,i]*wave[m,1,i]
            apdq[m,i] = s[2,i]*wave[m,2,i]
        end
    end

end #subroutine rp1