# ManagePic.jl

using FileIO, Images, ImageShow

"""
Convert a file to color matrices.

# Arguments
- `filename`::String`: name of file to load.
"""
function generateMatricesRGB(filename)
    img = load(filename)
    r = (Float64.(red.(img)))
    g = (Float64.(green.(img)))
    b = (Float64.(blue.(img)))
    return [r,g,b]
end # RGB

"""
Convert a loaded image to color matrices.

# Arguments
- `img`::Array{RGB{Normed{UInt8,8}},2}`: raw image.
"""
function generateMatricesRGB(img)
    r = (Float64.(red.(img)))
    g = (Float64.(green.(img)))
    b = (Float64.(blue.(img)))
    return [r,g,b]
end # RGB

"""
Convert color matrices o RGB matrice.

# Arguments
- `r`::Array{Float, 2}` full-size matrix of red color.
- `g`::Array{Float, 2}` full-size matrix of green color.
- `b`::Array{Float, 2}` full-size matrix of blue color.
"""
function matriceRGB(r, g, b)
    return ([RGB(r[i,j], g[i, j], b[i,j]) for i in 1:size(r)[1], j in 1:size(r)[2] ])
end # matriceRGB

"""
Save a r g b matrices as a picture.

# Arguments
- `newFilename`::String`: name of file.
- `r`::Array{Float, 2}` full-size matrix of red color.
- `g`::Array{Float, 2}` full-size matrix of green color.
- `b`::Array{Float, 2}` full-size matrix of blue color.
"""
function savePicture(newFilename, r, g, b)
    aa = matriceRGB(r, g, b)
    save(newFilename, aa)
end # save

"""
Save RGB matrice as a picture.

# Arguments
- `newFilename`::String`: name of file.
- `RGB`::Array{Array{Float64,2},1}`: RGB matrices list.
"""
function savePictureRGB(newFilename, RGB)
    savePicture(newFilename, RGB[1], RGB[2], RGB[3])
end # save2

"""
Convert RGB to HSL.
Return image as HSL matrices list.

# Arguments
- `rgb`::Array{Array{Float64,2},1}`: r, g, b matrices list.
"""
function rgb2hsl(rgb)
    r, g, b = rgb

    arraySize = size(r)
    h = zeros(Float64, arraySize[1], arraySize[2])
    s = zeros(Float64, arraySize[1], arraySize[2])
    l = zeros(Float64, arraySize[1], arraySize[2])

    # L:Lightness - value of white light
    l = lightness(rgb)

    Cmax = max.(r, g, b)
    Cmin = min.(r, g, b)
    Δ = Cmax .- Cmin

    for i in 1:arraySize[1], j in 1:arraySize[2]

        # H:Hue - value of color
        if Δ[i,j] == 0
            h[i,j] = 0
        elseif Cmax[i,j] == r[i,j]
            h[i,j] = 60 * (((g[i,j] - b[i,j]) ./ Δ[i,j]) .% 6)
        elseif Cmax[i,j] == g[i,j]
            h[i,j] = 60 * (((b[i,j] - r[i,j]) ./ Δ[i,j]) .+ 2)
        else
            h[i,j] = 60 * (((r[i,j] - g[i,j]) ./ Δ[i,j]) .+ 4)
        end

        # S:Saturation -  color saturation
        if Δ[i,j] == 0
            s[i,j] = 0
        else
            s[i,j] = Δ[i,j] ./ (1 - abs(2 * l[i,j] - 1))
        end
    end
    return [h,s,l]
end

"""
Convert HSL to RGB.
Return image as RGB matrices list.

#Arguments
- `hsl`::Array{Array{Float64,2},1}`: HSL matrices list.
"""
function hsl2rgb(hsl)
    h, s, l = hsl
    arraySize = size(h)
    r = zeros(Float64, arraySize[1], arraySize[2])
    g = zeros(Float64, arraySize[1], arraySize[2])
    b = zeros(Float64, arraySize[1], arraySize[2])

    for i in 1:arraySize[1], j in 1:arraySize[2]
        C = (1 - abs(2 * l[i,j] - 1)) .* s[i,j]
        X = C .* (1 - abs((h[i,j] ./ 60) % 2 - 1))
        m = l[i,j] - C ./ 2

        if h[i,j] >= 0 && h[i,j] < 60
            r[i,j] = C + m
            g[i,j] = X + m
            b[i,j] = m
        elseif h[i,j] >= 60 && h[i,j] < 120
            r[i,j] = X + m
            g[i,j] = C + m
            b[i,j] = m
        elseif h[i,j] >= 120 && h[i,j] < 180
            r[i,j] = m
            g[i,j] = m + C
            b[i,j] = m + X
        elseif h[i,j] >= 180 && h[i,j] < 240
            r[i,j] = m
            g[i,j] = X + m
            b[i,j] = C + m
        elseif h[i,j] >= 240 && h[i,j] < 300
            r[i,j] = X + m
            g[i,j] = m
            b[i,j] = C + m
        else
            r[i,j] = C + m
            g[i,j] = m
            b[i,j] = X + m
        end
    end
    return [r,g,b]
end

"""
Count lightness.

Arduments:
- `rgb`::Array{Array{Float64,2},1}`: r, g, b matrices list.
"""
function lightness(rgb)
    r, g, b = rgb
    l = (max.(r, g, b) .+ min.(r, g, b)) ./ 2
    return l
end


using Random, Distributions
using ProgressBars
using Plots

ENV["JULIA_DEBUG"] = "all"

const t = 3 # number of frames wrong
const expectedrates = [
    0.585, # delay-critical user u
    2, # delay-tolerant user v
]
const ζ = 2 # path loss exponent
const a = [
    0.7, # power for user u
    0.3,  # power for user v
]
const d = [
    2, # distance of user u
    1,  # distance of user v
]
const total_frames = 10000

channel_distribution = Normal(0, √2 / 2)
plot_points = ([], [])

Threads.@threads for SNR ∈ tqdm(0:5:45)
    for τ ∈ [0.6, 1]
        successful_frames = zeros((2, 1))

        progress_bar = 1:total_frames
        # progress_bar = ProgressBar(1:total_frames)
        # set_description(progress_bar, "($SNR dB, τ = $τ)")
        for frame ∈ progress_bar
            # Step 1: Calculate g of channel for each user
            h = rand(channel_distribution, (2, 2)) * [1, 1im]
            g = h ./ sqrt.(1 .+ d .^ ζ)

            # Step 1.1: Sort channel gains
            if abs2(h[1]) > abs2(h[2])
                h[2], h[1] = h[1], h[2]
                g = h ./ sqrt.(1 .+ d .^ ζ)
            end

            # Step 1.5: Apply time-correlated small fading coefficient
            ω = rand(channel_distribution, (2, 2)) * [1, 1im]
            h = τ^t * h + √(1 - τ^(2t)) * ω
            g = h ./ sqrt.(1 .+ d .^ ζ)

            # Step 2: Calculate AWGN
            σ = 10^(-SNR / 20)
            n = rand(Normal(0, σ), 2)

            # Step 3: Calculate reception SNR
            # Note: We are using the actual noise value n, instead of just the SNR-based σ
            SNRuu = a[1] * abs2(g[1]) / (a[2] * abs2(g[1]) + σ^2)
            SNRvv = a[2] * abs2(g[2]) / σ^2
            SNRvu = a[1] * abs2(g[2]) / (a[2] * abs2(g[2]) + σ^2)

            receiver_rates = [log2(1 + SNRuu), log2(1 + SNRvv), log2(1 + SNRvu)]

            successful_frames += [
                receiver_rates[1] >= expectedrates[1],
                receiver_rates[2] >= expectedrates[2] &&
                receiver_rates[3] >= expectedrates[1],
            ]
        end

        outage_rates = @. 1 - successful_frames / total_frames

        for rate ∈ outage_rates
            if rate > 0
                push!(plot_points[1], SNR)
                push!(plot_points[2], rate)
            end
        end
    end
end

## THEORETICAL PLOTS

# Calculate theoretical values
λ = 1 ./ (1 .+ d .^ ζ)
r = 2 .^ expectedrates .- 1
φ = max(r[1] / (a[1] - a[2] * r[1]), r[2] / a[2])

SNR = LinRange(0, 45, 100)
σ = @. 10^(-SNR / 20)
ρ = @. 1 / σ^2

plotly()
plot()
for τ ∈ [0.6, 1]
    # NOMA - Outdated CSI
    PoutTheoreticalNOMA = @. [
        1 - exp(-2 * r[1] * σ^2 / (2 - τ^(2t)) / (a[1] - a[2] * r[1]) / λ[1]),
        1 - 2 * exp(-φ * σ^2 / λ[2]) + exp(-2 * φ * σ^2 / (2 - τ^(2t)) / λ[2]),
    ]
    PoutTheoreticalOMA = @. [
        1 - exp(
            -2 * (2^(2 * expectedrates[1]) - 1) / ((2 - τ^(2t)) * ρ * λ[1]),
        ),
        1 - 2 * exp(-(2^(2 * expectedrates[2]) - 1) / ρ / λ[2]) + exp(
            -2 * (2^(2 * expectedrates[2]) - 1) / ((2 - τ^(2τ)) * ρ * λ[2]),
        ),
    ]

    label = [
        "a"
        "b"
    ]

    plot!(
        SNR,
        PoutTheoreticalNOMA,
        yscale = :log10,
        label = ["NOMA, User u, τ = $τ" "NOMA, User v, τ = $τ"],
        lw = 2
    )
    plot!(
        SNR,
        PoutTheoreticalOMA,
        yscale = :log10,
        label = ["OMA, User u, τ = $τ" "NOMA, User v, τ = $τ"],
        lw = 2
    )
end

plot!(
    plot_points[1],
    plot_points[2],
    seriestype = :scatter,
    yscale = :log10,
    label = "NOMA Simulation results",
)

xlabel!("SNR (dB)")
ylabel!("Outage rate")

display(plot!(legend = :best, size = (1000,800)))


using RecurrenceAnalysis, SparseArrays
using Test, Statistics, LinearAlgebra, DelimitedFiles

### Skeletonization of RPs
@testset "Skeletonized sinusoidal recurrence structures" begin
    data = sin.(2*π .* (0:1000)./ 60)
    Y = embed(data, 3, 15)

    RP = RecurrenceMatrix{FAN}(Y, 0.2)
    RP2 = RecurrenceMatrix(Y, 0.25; fixedrate=true)
    RP3 = RecurrenceMatrix(Y, 0.9)

    RP_skel = skeletonize(RP)
    RP_skel2= skeletonize(RP2)
    RP_skel3= skeletonize(RP3)

    @test RP_skel2 == RP_skel3

    tauRR1 = RecurrenceAnalysis.tau_recurrence(RP_skel)
    tauRR2 = RecurrenceAnalysis.tau_recurrence(RP_skel2)
    tauRR3 = RecurrenceAnalysis.tau_recurrence(RP_skel3)

    @test tauRR2 == tauRR3

    c1 = findall(x->x!=0, diff(tauRR1))
    c2 = findall(x->x!=0, diff(tauRR2))

    @test mean(diff(c1[3:2:end])) == 60
    @test mean(diff(c2[3:2:end])) == 60
end

@testset "Skeletonized chaotic recurrence structures" begin
    data = readdlm(joinpath(tsfolder, "test_time_series_lorenz_standard_N_10000_multivariate.csv"))

    RP = RecurrenceMatrix(Dataset(data[1:250,:]), 0.05; fixedrate=true)
    RP_skel = skeletonize(RP)

    @test RP_skel[135,61] == 1
    @test RP_skel[135,60] == 0
    @test RP_skel[135,60] == 0
    @test RP_skel[136,61] == 0
    @test RP_skel[134,61] == 0

    @test isempty(findall(!iszero,diag(RP_skel,1))) == true
    @test isempty(findall(!iszero,diag(RP_skel,-1))) == true

    d = Bool.(RP_skel[61:79,134:153])

    @test length(findall(diag(d,1))) == 18
    @test isempty(findall(diag(d))) == true
end


julia> q = Quaternion(1,0,0,0)
julia> q  = Quaternion (1, 2, 3, 4)
       q1 = Quaternion(2, 3, 4, 5)
       q2 = Quaternion(3, 4, 5, 6)
       r = 7.

julia> norm(q)
5.477225575051661

julia> -q
-1 - 2i - 3j - 4k

julia> conj(q)
1 - 2i - 3j - 4k

julia> r + q, q + r
(8.0 + 2.0i + 3.0j + 4.0k,8.0 + 2.0i + 3.0j + 4.0k)

julia> q1 + q2
5 + 7i + 9j + 11k

julia> r*q, q*r
(7.0 + 14.0i + 21.0j + 28.0k,7.0 + 14.0i + 21.0j + 28.0k)

julia> q1*q2, q2*q1, q1*q2 != q2*q1
(-56 + 16i + 24j + 26k,-56 + 18i + 20j + 28k,true)


@testset "bearings.jl" begin
    p1 = rad2deg(Point(0.5, 1.0))
    p2 = rad2deg(Point(1.0, 0.5))
    p3 = rad2deg(Point(0.6, 1.0))  # above p1

    rs = RouteSection(p1,p2)
    @test bearing(p1, p2) ≈ rad2deg(5.81412807071) atol = 0.001
    @test bearing(rs) ≈ rad2deg(5.81412807071) atol = 0.001
    @test bearing(p3, p1) ≈ 180.0 atol = 0.001

    @test final_bearing(p1, p2) ≈ rad2deg(5.45864813531) atol = 0.001
    @test final_bearing(rs) ≈ rad2deg(5.45864813531) atol = 0.001
end


using WeaklyHard
"""
Example for generating satisfaction sets for weakly-hard constraints
"""

# Constraints
lambda1 = AnyHitConstraint(12, 15)
lambda2 = RowHitConstraint(15, 47)
lambda3 = RowMissConstraint(3)

# Generate automata
G1 = build_automaton(lambda1)
G2 = build_automaton(lambda2)
G3 = build_automaton(lambda3)

# Generate satisfaction set of length N
N = 20
S1 = all_sequences(G1, N)
S2 = all_sequences(G2, N)
S3 = all_sequences(G3, N)

# check that all sequences satisfy the constraint
for w in S1
    if !is_satisfied(lambda1, w)
        println("Sequence ($(bitstring(w)[end-N+1:end])) does not satisfy $lambda1")
    end # if
end # for
for w in S2
    if !is_satisfied(lambda2, w)
        println("Sequence ($(bitstring(w)[end-N+1:end])) does not satisfy $lambda2")
    end # if
end # for
for w in S3
    if !is_satisfied(lambda3, w)
        println("Sequence ($(bitstring(w)[end-N+1:end])) does not satisfy $lambda3")
    end # if
end # for


# form intervals
function endPoints( A, V, M )
    nvec = size(V,2);
    p = kpm(A, V, M);

    ∂p = Polynomials.derivative(p);
    ∂2p = Polynomials.derivative(∂p);

    r_∂p = Polynomials.roots(∂p) |>
        x -> filter( isreal, x ) |>
        real                     |>
        x-> filter(
            item -> -1 < item < 1,
            x )

    r_∂2p = Polynomials.roots(∂2p) |>
        x -> filter( isreal, x )   |>
        real                       |>
        x-> filter(
            item -> -1 < item < 1,
            x )

    roots =  [ r_∂p; r_∂2p]
    roots,p
end


const data = parse.(Int, split(readline(), ","))

# Part One

minimum(UnitRange(extrema(data)...)) do i
    sum(abs.(data .- i))
end |> println

# Part Two

minimum(UnitRange(extrema(data)...)) do i
    sum(data) do d
        abs(d - i) * (abs(d - i) + 1) ÷ 2
    end
end |> println


#arc diagrams

# plots
#import Pkg
#Pkg.add("LinearAlgebra")
#Pkg.add("GraphRecipes")
#Pkg.add("SparseArrays")

using LinearAlgebra
using SparseArrays
using GraphRecipes
using Plots

adjmat = Symmetric(sparse(rand(0:1,8,8)))

plot(
    graphplot(adjmat,
              method=:arcdiagram,
              markersize=0.5,
              linecolor=:black,
              markercolor=:black)
    )


function sc(x)
  if(abs(x)>1.0)
	a=trunc(log10(abs(x)))
	b=x/(10.0^a)
  else
	a=trunc(log10(abs(x)))-1
	b=x/(10.0^a)
  end
  return a,b
end


using SingularSpectrumAnalysis
using Test, Statistics, Random
const SSA = SingularSpectrumAnalysis

@testset "SingularSpectrumAnalysis" begin
    Random.seed!(1)
    # Generate some data
    L = 20
    K = 100
    N = K*L;
    t = 1:N;
    T = 20;
    y = sin.(2pi/T*t);
    y .+= (0.5sin.(2pi/T*4*t)).^2
    e = 0.1randn(N);
    yn = y+e;
    # plot(yn)

    yt, ys = analyze(yn, L)
    A,x = fit_trend(yt, 1)
    @test x ≈ [0.124, 0] atol=0.1

    USV = hsvd(yn,L)

    trend, seasonal_groupings = autogroup(USV)
    @test trend == 3
    @test seasonal_groupings == [1:2, 4:5]

    yrt, yrs = reconstruct(USV, trend, seasonal_groupings)
    yr = sum([yrt yrs],dims=2)
    @test sqrt(mean((y.-yn).^2)) > sqrt(mean((y.-yr).^2))

    using Plots
    plot(USV, cumulative=true)
    plot(USV, cumulative=false)
    pairplot(USV, seasonal_groupings)
end
@testset "forecasting" begin
    include("forecast.jl")
end


using Plots
plotlyjs();

include("../julia/fitgamma.jl");
include("../julia/utils.jl");

x_true = readcsv("../data/spikes.csv")[:];
isi = diff(x_true);

lambda, k = fitgamma(isi);

@printf "Estimated parameters:\n"
@printf "lambda: %.4f, k: %.4f\n" lambda k


include("macros.jl")

Ad = [
   [-1.0/6.0  3.0/6.0 -3.0/6.0 1.0/6.0];
   [ 3.0/6.0 -6.0/6.0  0.0/6.0 4.0/6.0];
   [-3.0/6.0  3.0/6.0  3.0/6.0 1.0/6.0];
   [ 1.0/6.0  0.0/6.0  0.0/6.0 0.0/6.0]
]

dAd = [
   [ 0.0 -0.5  1.0 -0.5];
   [ 0.0  1.5 -2.0  0.0];
   [ 0.0 -1.5  1.0  0.5];
   [ 0.0  0.5  0.0  0.0]
]

d2Ad = [
   [ 0.0 0.0 -1.0  1.0];
   [ 0.0 0.0  3.0 -2.0];
   [ 0.0 0.0 -3.0  1.0];
   [ 0.0 0.0  1.0  0.0]
]

function eval_UC_spline(smin, smax, orders, C, S)

    d = size(S,2)
    N = size(S,1)

    vals = zeros(N)
    eval_UC_spline!(smin, smax, orders, C, S, vals)
    return vals

end

function eval_UC_spline_G(a, b, orders, C, S)

    d = size(S,2)
    N = size(S,1)

    vals = zeros(N)
    grad = zeros(N,d)

    eval_UC_spline_G!(a, b, orders, C, S, vals, grad)

    return (vals, grad)

end

function eval_UC_multi_spline(a, b, orders, C, S)

    d = size(S,2)
    N = size(S,1)
    K = size(C,1) # number of splines to evaluate
    vals = zeros(K,N)
    eval_UC_multi_spline!(a, b, orders, C, S, vals)
    return vals

end

# problem with this approach: the functions don't get cached.

# fun = (create_function(1,"natural"))
# fun.args[2].args[2]
# fun.args[2].args[2]
#
# for d = 1:4
#     eval(create_function(d,"natural"))
# end
#
# for d = 1:4
#     eval(create_function_with_gradient(d,"natural"))
# end

# with this approach functions don't get cached either (yet)
# but it's cleaner

@generated function eval_UC_spline!(a, b, orders, C, S, V)
    d = C.parameters[2]
    # the speed penalty of extrapolating points when iterating over point
    # seems very small so this is the default
    fun = (create_function(d,"natural"))
    return fun.args[2].args[2]
end

@generated function eval_UC_spline_G!(a, b, orders, C, S, V, dV)
    d = C.parameters[2]
    fun = (create_function_with_gradient(d,"natural"))
    return fun.args[2].args[2]
end

@generated function eval_UC_multi_spline!(a, b, orders, C, S, V)
    d = C.parameters[2]-1 # first dimension of C indexes the splines
    # the speed penalty of extrapolating points when iterating over point
    # seems very small so this is the default
    fun = (create_function_multi_spline(d,"natural"))
    return fun.args[2].args[2]
end


"""
    Parabolic{T <: Number} <: Construct{T}

A parabolic Cayley-Dickson construct as a pair of `Number`s.
"""
struct Parabolic{T <: Number} <: Construct{T}
    l::T
    r::T
end

function (==)(x::Parabolic, y::Parabolic)
    x.l == y.l && x.r == y.r
end

function real(z::Parabolic)
    real(z.l)
end

"""
    quadrance(z::CayleyDickson.Parabolic)

For parabolic Cayley-Dickson constructs, return the squared absolute value of the
left component of `z`. Depending on the type of the components of `z`, this can
either be non-negative, or positive, negative or zero.
"""
function quadrance(z::Parabolic)
    quadrance(z.l)
end

function iszerodivisor(z::Parabolic)
    iszerodivisor(z.l)
end

function zero(::Type{Parabolic{T}}) where T <: Number
    Parabolic{T}(zero(T), zero(T))
end

function zero(::Parabolic{T}) where T <: Number
    zero(Parabolic{T})
end

function one(::Type{Parabolic{T}}) where T <: Number
    Parabolic{T}(one(T), zero(T))
end

function one(::Parabolic{T}) where T <: Number
    one(Parabolic{T})
end

function (-)(z::Parabolic)
    Parabolic(-(z.l), -(z.r))
end

function conj(z::Parabolic)
    Parabolic(conj(z.l), -(z.r))
end

function (+)(x::Parabolic{T}, y::Parabolic{T}) where T <: Number
    Parabolic{T}(x.l + y.l, x.r + y.r)
end

function (+)(a::T, z::Parabolic{T}) where T <: Number
    Parabolic{T}(a + z.l, z.r)
end

function (+)(z::Parabolic{T}, a::T) where T <: Number
    Parabolic{T}(z.l + a, z.r)
end

function (+)(a::Real, z::Parabolic)
    l = a + z.l
    Parabolic(l, oftype(l, z.r))
end

function (+)(z::Parabolic, a::Real)
    l = z.l + a
    Parabolic(l, oftype(l, z.r))
end

function (-)(x::Parabolic{T}, y::Parabolic{T}) where T <: Number
    Parabolic{T}(x.l - y.l, x.r - y.r)
end

function (-)(a::T, z::Parabolic{T}) where T <: Number
    Parabolic{T}(a - z.l, -(z.r))
end

function (-)(z::Parabolic{T}, a::T) where T <: Number
    Parabolic{T}(z.l - a, z.r)
end

function (-)(a::Real, z::Parabolic)
    l = a - z.l
    Parabolic(l, -(oftype(l, z.r)))
end

function (-)(z::Parabolic, a::Real)
    l = z.l - a
    Parabolic(l, oftype(l, z.r))
end

function (*)(x::Parabolic{T}, y::Parabolic{T}) where T <: Number
    Parabolic{T}(
        (x.l * y.l),
        (y.r * x.l) + (x.r * conj(y.l))
    )
end

function (*)(a::T, z::Parabolic{T}) where T <: Number
    Parabolic{T}(a * z.l, z.r * a)
end

function (*)(z::Parabolic{T}, a::T) where T <: Number
    Parabolic{T}(z.l * a, z.r * conj(a))
end

function (*)(a::Real, z::Parabolic)
    Parabolic(a * z.l, a * z.r)
end

function (*)(z::Parabolic, a::Real)
    Parabolic(z.l * a, z.r * a)
end

function (/)(z::Parabolic, a::Real)
    if iszero(a)
        error(ZeroDenominator)
    end

    Parabolic(z.l / a, z.r / a)
end

function (\)(a::Real, z::Parabolic)
    if iszero(a)
        error(ZeroDenominator)
    end

    Parabolic(a \ z.l, a \ z.r)
end

function widen(::Type{Parabolic{T}}) where T
    Parabolic{widen(T)}
end

function float(::Type{Parabolic{T}}) where T <: AbstractFloat
    Parabolic{T}
end

function float(::Type{Parabolic{T}}) where T
    Parabolic{float(T)}
end

function cloak(z::Parabolic)
    Parabolic(cloak(z.l), dagger(z.r))
end

function dagger(z::Parabolic)
    Parabolic(dagger(z.l), cloak(z.r))
end

function star(z::Parabolic)
    Parabolic(dagger(star(z.r)), star(z.l))
end


export continuity

Space(d::IntervalOrSegment) = Chebyshev(d)
Space(d::FullSpace{<:Real}) = Chebyshev(Line())

Fun(::typeof(identity), d::IntervalOrSegment{T}) where {T<:Number} =
    Fun(Chebyshev(d), [mean(d), complexlength(d)/2])

## Calculus

# the default domain space is higher to avoid negative ultraspherical spaces
Integral(d::IntervalOrSegment,n::Integer) = Integral(Ultraspherical(1,d),n)


using Test
using OpenQuantumSystems
using Random, SparseArrays, LinearAlgebra

import OrdinaryDiffEq

@testset "schroedinger" begin

    Random.seed!(0)

    D(op1::Array, op2::Array) = abs(norm(op1 - op2))
    D(x1::StateVector, x2::StateVector) = norm(x2 - x1)
    D(op1::AbstractOperator, op2::AbstractOperator) =
        abs(tracedistance_nh(dense(op1), dense(op2)))
    D(op1::AbstractSuperOperator, op2::AbstractSuperOperator) =
        abs(tracedistance_nh(dense(op1), dense(op2)))

    N = 3
    Ncutoff = 2
    tspan = [0.0:0.1:1.0;]

    Ω = [
        0 2 3
        2 0 1
        3 1 0
    ]

    ω = [1.0, 1.2, 1.5]

    basis_fock = FockBasis(Ncutoff)
    basis = tensor([basis_fock for i = 1:N]...)

    a = destroy(basis_fock)
    at = create(basis_fock)
    I = identityoperator(basis_fock)

    psi0 = tensor([coherentstate(basis_fock, i % Ncutoff) for i = 1:N]...)

    # Interaction picture
    Hrot = SparseOpType[]
    for i = 1:N, j = 1:N
        if i == j
            continue
        end
        h = embed(basis, [i, j], [a, Ω[i, j] * at])
        push!(Hrot, h)
    end
    Hrot = sum(Hrot)

    # Schroedinger picture
    function f(t, psi)
        H = SparseOpType[embed(basis, i, ω[i] * at * a) for i = 1:N]
        for i = 1:N, j = 1:N
            if i == j
                continue
            end
            h = embed(basis, [i, j], [a, exp(1im * (ω[i] - ω[j]) * t) * Ω[i, j] * at])
            push!(H, h)
        end
        sum(H)
    end

    tout, psi_rot_t = schroedinger(psi0, tspan, Hrot)
    tout, psi_t = schroedinger_dynamic(psi0, tspan, f)

    n_op = dense(at * a)
    for (i, t) in enumerate(tout)
        R = prod([embed(basis, i, exp(1im * ω[i] * t * n_op)) for i = 1:N])
        psi_rot = psi_rot_t[i]
        psi = psi_t[i]
        # @test abs(dagger(psi_rot)*R*psi) < 1e-5
        rho = dm(psi)
        rho_rot = dm(psi_rot)
        @test tracedistance(rho_rot, dense(R) * rho * dagger(dense(R))) < 1e-5
    end

    function fout(t, psi)
        deepcopy(psi)
    end
    t_fout, psi_fout = schroedinger(psi0, tspan, Hrot; fout = fout)
    @test t_fout == tout && psi_fout == psi_rot_t

    # test integration of propagator using 2 level system
    basis = SpinBasis(1 // 2)
    su = spinup(basis)
    u0 = identityoperator(basis)
    sx = sigmax(basis)
    sz = sigmaz(basis)

    # for the time dependent equation
    f2(t, psi) = sx * π
    tspan = [0:1.0;]
    t, u = schroedinger(u0, tspan, π * sx)

    # I think the tolerance on the differential equation is 1e-6, we expect the operator to be essentially the identity
    @test abs(expect(sz, u[end] * su)) - abs(expect(sz, u0 * su)) < 1e-6

    t, u = schroedinger_dynamic(u0, tspan, f2)
    @test abs(expect(sz, u[end] * su)) - abs(expect(sz, u0 * su)) < 1e-6

    mode1 = Mode(0.2, 1.0)
    Energy = [0.0, 200.0]
    mol1 = Molecule([mode1], 2, Energy)
    mol2 = Molecule([mode1], 2, Energy)
    agg = Aggregate([mol1, mol2])
    aggInds = getIndices(agg)
    aggIndsLen = length(aggInds)
    basis = GenericBasis([aggIndsLen])
    FCFact = getFranckCondonFactors(agg, aggInds)
    Ham = getAggHamiltonian(agg, aggInds, FCFact)

    ket0 = randstate(basis)
    ket0 = Ket(basis)
    for i = 1:aggIndsLen
        ket0.data[i] = (1.0 / aggIndsLen)^0.5
    end
    rho0 = dm(ket0)
    tspan = [0.0:0.1:1.0;]

    T, rho_t = schroedinger(
        ket0,
        tspan,
        Ham;
        reltol = 1e-10,
        abstol = 1e-10,
        alg = OrdinaryDiffEq.Tsit5(),
    )
    for t_i = 1:length(tspan)
        U_op = evolutionOperator(Ham, tspan[t_i])
        ket = U_op * ket0
        @test 1e-8 > D(ket, rho_t[t_i])
        # println(t_i, " ", D(ket.data, rho_t[t_i].data))
    end

end # testset


function eval_overlapᵥ(
    xₜ::Vector{Float64},
    x::Matrix{Float64},
    fs_counts::Vector{Int64},
    ol_maxₜ::Vector{Float64},
    v::Int64,
)
    xₜ = xₜ .* fs_counts
    cons = LinearAlgebra.BLAS.gemv('T', x, xₜ) - ol_maxₜ
    cons[v] = 0
    return sum_pos(cons)::Float64
end

function eval_overlap(
    x::Matrix{Float64},
    fs_counts::Matrix{Float64},
    ol_max::Matrix{Float64},
    T::Int64,
    ol2::Vector{Float64},
)
    fs_counts = fs_counts .* x
    ol = _gemmblas(x, fs_counts, ol_max)
    for v2 = 1:T
        olₜ = ol[:, v2]
        olₜ[v2] = zero(Float64)
        ol2[v2] = sum_pos(olₜ)
    end
    return ol2::Vector{Float64}
end

function eval_exp_scoreₜ(x_Iₜ::Vector{Float64}, expected_scoreₜ::ExpectedScore) #x = I
    es = (expected_scoreₜ.val * x_Iₜ) ./ sum(x_Iₜ)
    # _gemvblasT(
    #     expected_scoreₜ.val,
    #     x_Iₜ,
    #     size(expected_scoreₜ.val, 2)
    # ) ./ sum(x_Iₜ)
    es = vcat((es .- expected_scoreₜ.max)..., (-es .+ expected_scoreₜ.min)...)
    return es
end

function eval_cons(xₜ::Vector{Float64}, Aₜ::Matrix{Float64}, bₜ::Vector{Float64})
    cons = copy(bₜ)
    cons = _gemvblas(Aₜ, xₜ, cons, size(xₜ, 1))
    return cons::Vector{Float64}
end

function check_feas(
    fs::FriendSets,
    consₜ::Constraint,
    xₜ::Vector{Float64},
    n_fs::Int64,
    n_items::Int64,
)
    es = 0
    cons = 0
    if size(consₜ.constr_A, 1) > 0
        cons = eval_cons(xₜ, consₜ.constr_A, consₜ.constr_b)
    end
    if n_items > n_fs
        x_Iₜ = fs_to_items(xₜ, n_items, fs.items)
    else
        x_Iₜ = copy(xₜ)
    end
    if size(consₜ.expected_score.val, 1) > 0
        cons = vcat(cons, eval_exp_scoreₜ(x_Iₜ, consₜ.expected_score))
    end
    return sum_pos(cons)::Float64, x_Iₜ::Vector{Float64}
end


#
"""
 Functional based mesh adaptation

 FWD Pass: NonLin layer -> least sq. on linear layer
 BWD Pass: GD all the way through (takes adjoint on least sq.)

 Task: replace least sq. with PDE solve
       test problem: steady advection diffusion eqn at large
                     Peclet # (thin thin boundary layer)

"""

using Revise
using diffMesh
import Zygote
using NNlib, Flux, LinearAlgebra, Plots
#--------------------------------------#
np  = 2001 # data points
ng  = 31   # FEM nodes

function udata(x)
    u   = 2 .+ zero(x)
    u .+= @. sin(1*pi*x) * (-1 <= x <= 0)
    u .+= @. sin(4*pi*x) * ( 0 <  x <= 1)
    return u
end
#--------------------------------------#
function addEndPts(xe::AbstractArray
                  ,x0::Number,x1::Number)

    n = length(xe) + 2
    A,R = Zygote.ignore() do
        A = zeros(n,2)
        A[1,1] = 1
        A[n,2] = 1
        R = restrict(n)
        return A,R
    end

    xg  = A * [x0,x1] + R'*xe

    return xg
end
#--------------------------------------#
function model(xe,xp,up) # grid points, eval points, values to fit

    xg = addEndPts(xe,-1,1)
    NN = reluHatFunc(xg)
    ug = llsq(NN,xp,up)

    return reluHat(NN,ug,xp)
end
#--------------------------------------#
function llsq(NN,xp,up)

    A  = NN(xp)
    linOp(u) = reluHat(NN,u,xp)

#   u = udata(x)
    u = A \ up
    return u
end
#--------------------------------------#
function loss()
    fp   = model(xe,xp,up)
    loss = Flux.Losses.mse(fp,up)
    return loss
end
#--------------------------------------#
function makeplot()
    fp = model(xe,xp,up)

    xg = addEndPts(xe,-1,1)
    NN = reluHatFunc(xg)
    ug = llsq(NN,xp,up)

    plt = plot()
    plt = plot!(title="FEM AutoMesh",xlabel="x",ylabel="y")

    plt = plot!(xp,up,width=3,label="data")
    plt = plot!(xp,fp,width=3,label="prediction")
    plt = scatter!(xg,ug,width=4,color="red",label=false)

    # hat functions
    NNp = NN(xp')
    for i=1:ng
        scale = ug[i]
#       scale = 1.0
        global plt = plot!(xp,scale .* NNp[i,:]
                          ,width=2,color="black",label=false)
    end

    #savefig(plt,"diffMesh.png")
    display(plt)
    return plt
end
#--------------------------------------#
function callback(;doplot=true)

    global xe .= sort(xe)
    l2 = loss()
    fp = model(xe,xp,up)
    er = norm(up .- fp,Inf)
    println("Error ∞ norm: $er, loss: $l2")

    if(doplot) makeplot() end
    return false
end
#--------------------------------------#

# data
xp = Array(range(-1,stop=1,length=np))'
up = udata(xp)

# initialize
xg = Array(range(-1,stop=1,length=ng))
xg = sort(2*rand(ng).-1)

xe = xg[2:end-1]
fp = model(xe,xp,up)

callback(doplot=true)

ps = Flux.params(xe)
data = Iterators.repeated((),100)
opt = ADAM(1e-2)

gs = gradient(()->loss(),ps)

#anim = Animation()
for d in data
    global gs = gradient(()->loss(),ps)
    Flux.Optimise.update!(opt,ps,gs)
    callback(doplot=true)
    #frame(anim)
end
#gif(anim,"examples/AutoMesh.gif",fps=30)

#--------------------------------------#
nothing


using Statistics

include("update.jl")

function randomwalk(num::Integer, nsteps::Integer, prob::Real=0.5)
    means, vars = zeros(nsteps+1), zeros(nsteps+1)

    xs = zeros(num)
    means[1] = mean(xs)
    vars[1] = var(xs)
    for t in 1:nsteps
        map!(x -> update(x, prob), xs, xs)
        means[t+1] = mean(xs)
        vars[t+1] = var(xs)
    end
    return means, vars
end


#
# Potential energy and force
#
function forcepair(x,y,r,data)
  r6 = r^6
  r12 = r6^2
  r7 = r6*r
  r13 = r12*r

  drdx1 = -(x[1]-y[1])/r
  drdx2 = -(x[2]-y[2])/r

  sigfac1 = -12*(data.sig12/r13)
  sigfac2 = -6*(data.sig6/r7)

  dfacdr = -data.eps4*(sigfac1 - sigfac2)

  upair = data.eps4*(data.sig12/r12 - data.sig6/r6)
  fx = ( -dfacdr*drdx1, -dfacdr*drdx2 )

  return upair, fx
end


function \(cur_matrix::FourBlockMatrix, cur_vector::Vector)

  A, B, C, D = get_sub_matrices(cur_matrix)

  iszero(cur_matrix.n) && return D \ cur_vector
  iszero(cur_matrix.m) && return A \ cur_vector

  top_b = view(cur_vector, 1:cur_matrix.n)
  bot_b = view(cur_vector, cur_matrix.n .+ (1:cur_matrix.m) )

  cur_lu = lu(A)
  F = SharedArray{Float64}(cur_matrix.m,cur_matrix.m)

  @sync @distributed for j = 1:size(B,2)
    tmp_B = B[:,j]
    isempty(tmp_B) && continue
    F[:,j] = C * ( cur_lu \ Vector(tmp_B) )
  end

  G = D - F

  L_1 = A \ top_b

  L_2 = G \ ( C * L_1 )

  top_left = L_1 + A \ ( B * L_2 )
  bot_left = - L_2

  R_1 = G \ bot_b

  top_right = - A \ ( B * R_1 )
  bot_right = R_1

  vcat(
    ( top_left .+ top_right ),
    ( bot_left .+ bot_right )
  )

end


module YLEFlow

using Symbolics

# Write your package code here.

@variables t k ρ a0 a1 a2 a3 a4 a5 a6 a7 a8
@variables U(..) af1(..) af2(..) af3(..) af4(..) af5(..) af6(..) af7(..) af8(..)

msigma2 = 2*ρ*Differential(ρ)(Differential(ρ)(U(ρ))) + Differential(ρ)(U(ρ))
#dUdk = exp(5*t)/(exp(2*t) + msigma2)
dUdk = k^5/(k^2 + msigma2)

S1 =  [U(ρ),af1(ρ),af2(ρ),af3(ρ),af4(ρ),af5(ρ),af6(ρ),af7(ρ),af8(ρ)]
S2 =  [af1(ρ),af2(ρ),af3(ρ),af4(ρ),af5(ρ),af6(ρ),af7(ρ),af8(ρ),0]
S3 = [a0, a1, a2, a3, a4, a5, a6, a7, a8]

R=[ρ,ρ,ρ,ρ,ρ,ρ,ρ,ρ,ρ]
R[1]=dUdk

Q=similar(R)
for (c,p) in zip(S1,S2)
    R[1]=simplify(substitute(R[1],Dict( Differential(ρ)(c)  => p)))
end

for i in 2:length(R)
    R[i] = Symbolics.derivative(R[i-1],ρ)

    for (c,p) in zip(S1,S2)
        R[i]=(substitute(R[i],Dict( Differential(ρ)(c)  => p)))
    end
end

Q = copy(R)

for (c,p) in zip(S1,S3)
    for i in 1:length(R)
        Q[i]=substitute(Q[i],Dict( c  => p))
    end
end

for i in 1:length(R)
    Q[i]=simplify(Q[i])
end

Symbolics.derivative(Q[2],a2)

dU1dk = Symbolics.build_function(Q[2],k,a1,a2,a3,a4,a5,a6,a7,a8,ρ,expression=Val{false})
dU2dk = Symbolics.build_function(Q[3],k,a1,a2,a3,a4,a5,a6,a7,a8,ρ,expression=Val{false})
dU3dk = Symbolics.build_function(Q[4],k,a1,a2,a3,a4,a5,a6,a7,a8,ρ,expression=Val{false})
dU4dk = Symbolics.build_function(Q[5],k,a1,a2,a3,a4,a5,a6,a7,a8,ρ,expression=Val{false})
dU5dk = Symbolics.build_function(Q[6],k,a1,a2,a3,a4,a5,a6,a7,a8,ρ,expression=Val{false})
dU6dk = Symbolics.build_function(Q[7],k,a1,a2,a3,a4,a5,a6,a7,a8,ρ,expression=Val{false})
dU7dk = Symbolics.build_function(Q[8],k,a1,a2,a3,a4,a5,a6,a7,a8,ρ,expression=Val{false})
dU8dk = Symbolics.build_function(Q[9],k,a1,a2,a3,a4,a5,a6,a7,a8,ρ,expression=Val{false})

function drodk(dU1dk,dU2dk,a2,a3,ρ)
    return  -(dU1dk + 2*ρ*dU2dk)/(3*a2+2*ρ*a3)
end

function flow!(da,a,m2,t)
 k = exp(t)
 ρ = a[1]
 a1 = m2 - 2ρ * a[2]

 dU1dk_s = dU1dk(t, a1, a[2], a[3], a[4], a[5], a[6], a[7], a[8], ρ)
 dU2dk_s = dU2dk(t, a1, a[2], a[3], a[4], a[5], a[6], a[7], a[8], ρ)
 dU3dk_s = dU3dk(t, a1, a[2], a[3], a[4], a[5], a[6], a[7], a[8], ρ)
 dU4dk_s = dU4dk(t, a1, a[2], a[3], a[4], a[5], a[6], a[7], a[8], ρ)
 dU5dk_s = dU5dk(t, a1, a[2], a[3], a[4], a[5], a[6], a[7], a[8], ρ)
 dU6dk_s = dU6dk(t, a1, a[2], a[3], a[4], a[5], a[6], a[7], a[8], ρ)
 dU7dk_s = dU7dk(t, a1, a[2], a[3], a[4], a[5], a[6], a[7], a[8], ρ)
 dU8dk_s = dU8dk(t, a1, a[2], a[3], a[4], a[5], a[6], a[7], a[8], ρ)

 da[1] = drodk(dU1dk_s,dU2dk_s,a[2],a[3],ρ)

 da[2] = dU2dk_s + a[3] * da[1]
 da[3] = dU3dk_s + a[4] * da[1]
 da[4] = dU4dk_s + a[5] * da[1]
 da[5] = dU5dk_s + a[6] * da[1]
 da[6] = dU6dk_s + a[7] * da[1]
 da[7] = dU7dk_s + a[8] * da[1]
 da[8] = dU8dk_s
end

function g()

end

end


@system ThermalTime begin
    T: temperature ~ track(u"°C", override)
    Δt(context.clock.step): timestep ~ preserve(u"hr")
    ΔT: magnitude ~ track
    r(ΔT, Δt): rate => ΔT / Δt ~ track(u"hr^-1")
end

@system GrowingDegree(ThermalTime) begin
    Tb: base_temperature ~ preserve(u"°C", extern, parameter)
    To: optimal_temperature ~ preserve(u"°C", optional, extern, parameter)
    Tx: maximum_temperature ~ preserve(u"°C", optional, extern, parameter)

    ΔT(T, Tb, To, Tx): magnitude => begin
        T = !isnothing(To) ? min(T, To) : T
        T = !isnothing(Tx) && T >= Tx ? Tb : T
        T - Tb
    end ~ track(u"K", min=0)
    r(ΔT, Δt): rate => ΔT / Δt ~ track(u"K/hr")
end

@system BetaFunction(ThermalTime) begin
    Tn: minimum_temperature => 0 ~ preserve(u"°C", extern, parameter)
    To: optimal_temperature ~ preserve(u"°C", extern, parameter)
    Tx: maximum_temperature ~ preserve(u"°C", extern, parameter)
    β: beta => 1 ~ preserve(parameter)

    ΔT(T, Tn, To, Tx, β): magnitude => begin
        # beta function, See Yin et al. (1995), Ag For Meteorol., Yan and Hunt (1999) AnnBot, SK
        if (Tn < T < Tx) && (Tn < To < Tx)
            Ton = To - Tn
            Txo = Tx - To
            f = (T - Tn) / Ton
            g = (Tx - T) / Txo
            α = β * (Ton / Txo)
            f^α * g^β
        else
            0.
        end
    end ~ track
end

@system Q10Function(ThermalTime) begin
    To: optimal_temperature ~ preserve(u"°C", extern, parameter)
    Q10 => 2 ~ preserve(extern, parameter)

    #FIXME: Q10 isn't actually a thermal function like others (not a rate, check unit)
    ΔT(T, To, Q10): magnitude => begin
        Q10^((T - To) / 10u"K")
    end ~ track
end

export GrowingDegree, BetaFunction, Q10Function


@testset "Interface" begin
    rng = Random.MersenneTwister(42)
    D = 5
    prior = MvNormal(0.5 * ones(D))
    likelihood = MvNormal(7.0 * ones(D))
    logintegrand = x -> logpdf(likelihood, x)
    m = BayesModel(prior, logintegrand)
    s = PriorSampling()
    bquad = BayesQuad(SEKernel() ∘ ScaleTransform(0.1))
    pZ, x = bquad(m, s; nsamples=100)
    trueZ = exp(-0.5 * (logdet(cov(prior) + cov(likelihood)) + D * log(2π)))
    @test mean(pZ) ≈ trueZ atol = 1e-3
end


module SchurPolynomials

export schur, semi_standard_young_tableaux

using AbstractAlgebra, Transducers, Dictionaries
using YAPP: Monomial, Polynomial
using BangBang
using AbstractAlgebra: Partition
using Transducers: next, complete, Foldable, Map

const RowsType = Vector{SubArray{Int, 1, Vector{Int}, Tuple{UnitRange{Int}}, true}}

"""
Iterates over all semi-standard Young tableaux of shape `λ`.
"""
struct SemiStandardYoungTableaux <: Foldable
    n::Int
    T::Generic.YoungTableau{Int}
    rows::RowsType
    function SemiStandardYoungTableaux(n, λ)
        fill = Vector{Int}(undef, sum(λ))
        rows = map(λ, Iterators.accumulate(+, λ)) do len, stop
            view(fill, stop-len+1:stop)
        end
        return new(n, YoungTableau(λ, fill), rows)
    end
end

"""
Iterates over all configurations of row `i` that are allowed by the previous row.
Iterator returns `nothing` but mutates `row`.
"""
struct RowConfigs <: Foldable
    rows::RowsType
    i::Int
    n::Int
end
Base.eltype(::RowConfigs) = Nothing
function Transducers.__foldl__(rf::RF, acc, (; rows, i, n)::RowConfigs) where {RF}
    row = @inbounds rows[i]
    acc = _row!(i, 1, row, rows, n, rf, acc)
    return complete(rf, acc)
end
# row i, column j
function _row!(i, j, row, rows, n, rf::RF, acc::A) where {RF, A}
    j > length(row) && return next(rf, acc, nothing)
    above = i == 1 ? 0 : rows[i-1][j]
    left = get(row, j-1, 1)
    @inbounds for row[j] in max(above + 1, left):n
        acc = _row!(i, j + 1, row, rows, n, rf, acc)
        acc isa Reduced && return acc
    end
    acc
end

function _worker!(i, rf::RF, acc, ssyt) where {RF}
    (; n, T, rows) = ssyt
    i > length(rows) && return next(rf, acc, nothing)
    foldl(RowConfigs(rows, i, n); init=acc) do acc, _
        _worker!(i + 1, rf, acc, ssyt)
    end
end

function Transducers.__foldl__(rf::RF, acc, ssyt::SemiStandardYoungTableaux) where {RF}
    acc = _worker!(1, rf, acc, ssyt)
    return complete(rf, acc)
end

function Base.length((; n, T)::SemiStandardYoungTableaux)
    return Int(prod(
        (n-i+j) // hooklength(T, i, j)
        for (i, j_max) in enumerate(T.part)
        for j in 1:j_max
    ))
end
Base.eltype(::SemiStandardYoungTableaux) = Generic.YoungTableau

semi_standard_young_tableaux(n, λ) = SemiStandardYoungTableaux(n, λ) |> Map() do T
    YoungTableau(T.part, copy(T.fill))
end

function schur(n, λ)
    SSYT = SemiStandardYoungTableaux(n, λ)
    (; T) = SSYT
    foldl(SSYT; init=Polynomial{Int}()) do p, _
        exponents = map(1:n) do i
            count(==(i), T.fill)
        end
        add!!(p, Monomial(exponents))
    end
end

end


using Test, SummationByPartsOperators

for T in (Float32, Float64), split_form in (Val{true}(), Val{false}())
    xmin = T(-1)
    xmax = T(1)
    N = 2^6
    u0func = sinpi
    tspan = (zero(T), one(T))
    afunc = one

    # Legendre
    let D = legendre_derivative_operator(xmin, xmax, N)
        Di = nothing
        semidisc = VariableLinearAdvectionNonperiodicSemidiscretisation(D, Di, afunc, split_form, zero, zero)
        ode = semidiscretise(u0func, semidisc, tspan)
        du = similar(ode.u0)
        semidisc(du, ode.u0, nothing, first(tspan))
    end

    # SBP FD
    for acc_order in 2:2:8
        D = derivative_operator(MattssonSvärdNordström2004(), 1, acc_order, xmin, xmax, N)
        Di = dissipation_operator(D)
        semidisc = VariableLinearAdvectionNonperiodicSemidiscretisation(D, Di, afunc, split_form, zero, zero)
        ode = semidiscretise(u0func, semidisc, tspan)
        du = similar(ode.u0)
        semidisc(du, ode.u0, nothing, first(tspan))
    end
end


include("ZDT.jl")
include("DTLZ.jl")

# artificial problem
function MTP(D = 10, n_solutions = 100)
    f(x) = begin
        Q = 1.0 .+ sum((x[3:end] .- 0.8).^2)
        fx = [x[1], x[2]] .* Q
        g = [ (x[1] - 1)^2 + (x[2] - 1)^2 - 1 ]
        h = [0.0]
        return fx, g, h
    end

    bounds = Array([ zeros(D) 2.0ones(D)]')
    θ = range(π, 1.5π, length=n_solutions)

    x = 1.0 .+ cos.(θ)
    y = 1.0 .+ sin.(θ)

    pareto_set = [ generateChild(zeros(0), ([x[i], y[i]], [0.0], [0.0])) for i in 1:length(θ) ]

    return f, bounds, pareto_set

end


include("Quad/Spaces.jl")
include("Quad/Genus.jl")
include("Quad/GenusRep.jl")
include("Quad/Zlattices.jl")


using BookBayes
using LightGraphs

C = Variable(:c, 3);
S = Variable(:s, 3);
V = Variable(:v, 3);

# P(C)
φC = Factor([C], FactorTable(
    (c=1,) => 0.80,
    (c=2,) => 0.19,
    (c=3,) => 0.01,
    ));
# P(S|C)
φSC = Factor([S,C], FactorTable(
    (s=1,c=1,) => 0.001,
    (s=1,c=2,) => 0.200,
    (s=1,c=3,) => 0.800,
    (s=2,c=1,) => 0.009,
    (s=2,c=2,) => 0.750,
    (s=2,c=3,) => 0.199,
    (s=3,c=1,) => 0.990,
    (s=3,c=2,) => 0.050,
    (s=3,c=3,) => 0.001,
    ));

# P(V|C)
φVC = Factor([V,C], FactorTable(
    (v=1,c=1,) => 0.2,
    (v=1,c=2,) => 0.5,
    (v=1,c=3,) => 0.4,
    (v=2,c=1,) => 0.2,
    (v=2,c=2,) => 0.4,
    (v=2,c=3,) => 0.4,
    (v=3,c=1,) => 0.6,
    (v=3,c=2,) => 0.1,
    (v=3,c=3,) => 0.2,
    ));

g = SimpleDiGraph(3)

add_edge!(g, 1, 2)
add_edge!(g, 1, 3);

bn = BayesianNetwork([C,S,V], [φC, φSC, φVC], g);

m_weighted = LikelihoodWeightedSampling(100000);
m_var_elim = VariableElimination(3:1);

r_exact= infer(ExactInference(), bn, (:c,), (s=2, v=1))
r_weighted = infer(m_weighted, bn, (:c,), (s=2,v=1,))
r_var_elim = infer(m_var_elim, bn, (:c,), (s=2,v=1,))


sz_half = [0.5 0.0; 0.0 -.5]
@testset "MPS" begin
    d, lx = 2, 8

    # A uniformly chosen random state as MPS
    ketstate = complex.(randn(d^lx), randn(d^lx))
    randmps = MPS(lx, d, ketstate)

    # The Bethe chain ground state as MPS
    H = xxz_hamiltonian(lx)
    eheis, vheis = eigsolve(H, 1, :SR)
    mps = MPS(lx, 2, vheis[1])

    @testset "ketstate to MPS to ketstate" begin
        @test norm(ketstate) ≈ norm(randmps)
        @test ketstate ≈ mps2ketstate(randmps)
    end

    @testset "overlap of two MPS" begin
        @test sqrt(overlap(mps, mps)) ≈ norm(mps)
        @test overlap(mps, randmps) ≈ conj.(overlap(randmps, mps))
    end

    @testset "operator measurements" begin
        sz = spinoperators(1/2)[1]
        szdata = measure(mps, sz)
        @test all([isapprox(szdata[i], 0.0, atol=1.e-12) for i=1:lx])

        mpo = xxz_mpo(Float64, lx, 2)
        @test measure(mps, mpo)[1] ≈ eheis[1]
    end

    @testset "MPS center manipulations" begin
        ### TODO: write more tests for here!
        center_at!(randmps, lx-1)
        center_at!(randmps, lx-2)
        center_at!(randmps, 1)
        center_at!(randmps, lx)
        @test mps2ketstate(randmps) ≈ ketstate

        center_at!(randmps, lx)
        @test mps2ketstate(randmps) ≈ ketstate
    end

end

@testset "apply" begin
    @testset "twosite" begin
        lx, l = 8, 3

        H = xxz_hamiltonian(lx)
        eheis, vheis = eigsolve(H, 1, :SR)
        mps = MPS(lx, 2, vheis[1])

        mpscopy = deepcopy(mps)
        center_at!(mps, l)
        sz = spinoperators(1/2)[1]
        M = reshape(kron(sz, sz), 2, 2, 2, 2)
        apply!(mps, M, l)

        #@test mps_dims_are_consistent(mps)
        @test overlap(mpscopy, mps) ≈ measure(mpscopy, sz, sz, l, l+1)
    end
end


export QuLDE, LDEMSAlgHHL, QuNLDE
export QuEuler, QuLeapfrog, QuAB2, QuAB3, QuAB4

abstract type QuODEAlgorithm <: DiffEqBase.AbstractODEAlgorithm end
"""
    LDEMSAlgHHL <: QuODEAlgorithm

Multi-step methods based on HHL

"""
abstract type LDEMSAlgHHL <: QuODEAlgorithm end

"""
    QuLDE <: QuODEAlgorithm

Linear differential equation solvers (non-HHL)
    * k : order of Taylor series expansion

"""
struct QuLDE <: QuODEAlgorithm
    k::Int
    QuLDE(k = 3) = new(k)
end

"""
    QuNLDE <: QuODEAlgorithm

Linear differential equation solvers (non-HHL)
    * k : order of Taylor series expansion
    * ϵ : precision

"""
struct QuNLDE <: QuODEAlgorithm
    k::Int
    ϵ::Real
    QuNLDE(k = 3, ϵ = 1e-3) = new(k,ϵ)
end

"""
    QuEuler{T} <: LDEMSAlgHHL

Euler Method using HHL (1-step method)

"""
struct QuEuler{T} <: LDEMSAlgHHL
    step::Int
    α::Vector{T}
    β::Vector{T}
    nreg::Int
    QuEuler(nreg = 12,::Type{T} = Float64) where {T} = new{T}(1,[1.0,],[1.0,],nreg)
end

"""
    QuLeapfrog{T} <: LDEMSAlgHHL

Leapfrog Method using HHL (2-step method)

"""
struct QuLeapfrog{T} <: LDEMSAlgHHL
    step::Int
    α::Vector{T}
    β::Vector{T}
    nreg::Int
    QuLeapfrog(nreg = 12,::Type{T} = Float64) where {T} = new{T}(2,[0, 1.0],[2.0, 0],nreg)
end

"""
    QuAB2{T} <: LDEMSAlgHHL

AB2 Method using HHL (2-step method)

"""
struct QuAB2{T} <: LDEMSAlgHHL
    step::Int
    α::Vector{T}
    β::Vector{T}
    nreg::Int
    QuAB2(nreg = 12,::Type{T} = Float64) where {T} = new{T}(2,[1.0, 0], [1.5, -0.5],nreg)
end

"""
    QuAB3{T} <: LDEMSAlgHHL

AB3 Method using HHL (3-step method)

"""
struct QuAB3{T} <: LDEMSAlgHHL
    step::Int
    α::Vector{T}
    β::Vector{T}
    nreg::Int
    QuAB3(nreg = 12,::Type{T} = Float64) where {T} = new{T}(3,[1.0, 0, 0], [23/12, -16/12, 5/12],nreg)
end

"""
    QuAB4{T} <: LDEMSAlgHHL

AB4 Method using HHL (4-step method)

"""
struct QuAB4{T} <: LDEMSAlgHHL
    step::Int
    α::Vector{T}
    β::Vector{T}
    nreg::Int
    QuAB4(nreg = 12,::Type{T} = Float64) where {T} = new{T}(4,[1.0, 0, 0, 0], [55/24, -59/24, 37/24, -9/24],nreg)
end


###########
# This gives an example of Gauss–Seidel
#
###########

using BlockBandedMatrices, BandedMatrices, ArrayLayouts, FillArrays

function finitedifference_2d(n)
    h = 1/n
    D² = BandedMatrix(0 => Fill(-2,n), 1 => Fill(1,n-1), -1 => Fill(1,n-1))/h^2
    D_xx = BandedBlockBandedMatrix(Kron(D², Eye(n)))
    D_yy = BandedBlockBandedMatrix(Kron(Eye(n), D²))
    D_xx + D_yy
end

function _gaussseidel(L, U, b, x=copy(b), y=copy(b), M=5)
    for _=1:M
        @view(y[1:end-1]) .= Mul(U , @view(x[2:end]))
        y[end] = 0
        x .= b .- y
        x .= Ldiv(L, x)
    end
    x
end

function _gaussseidel2(L, U, b, x=copy(b), y=copy(b), M=5)
    for _=1:M
        mul!(@view(y[1:end-1]) , U , @view(x[2:end]))
        y[end] = 0
        y .= b .- y
        x .= L\ y
    end
    x
end

function gaussseidel(A, b, M)
    n = Int(sqrt(length(b)))
    L = LowerTriangular(A)
    U = BandedBlockBandedMatrix(UpperTriangular(@view(A[1:end-1,2:end])), ([fill(n,n-1); n-1], [n-1; fill(n,n-1)]),
                                            (0,1), (0,1))
    x = copy(b)
    _gaussseidel(L,U, b, x, copy(x), M) # 1.6s
    x
end

n = 1000;
    Δt = (1/n^2)/4; @time Δ = finitedifference_2d(n);
    @time A = I - Δt*Δ  # 1m x 1m discretization
    b = randn(n^2)
    L = LowerTriangular(A)
    @time U = BandedBlockBandedMatrix(UpperTriangular(@view(A[1:end-1,2:end])), ([fill(n,n-1); n-1], [n-1; fill(n,n-1)]),
                                            (0,1), (0,1))
    x = copy(b)
    y = copy(x)
    @time _gaussseidel(L,U, b, x, y, 20) # 0.4s

    @time L̃ = sparse(BandedBlockBandedMatrix(L))
    @time Ũ = sparse(U)
    x = copy(b)
    y = copy(x)
    @time _gaussseidel2(L̃,Ũ, b, x, y, 20)

Ã = sparse(A)
@time qr(Ã);
@time Ã \b ;

A = randn(n,n)

n = 1000;
    As = [randn(2k+1,k) for k=1:n];

    @time for A in As
        qr(A)
    end
n = 1000;
    # As = [randn(2k+1,k) for k=1:n];

    @time for A in As
        lu(A)
    end

0.02*n

A*x - b

x̃ = copy(x)
ỹ = copy(y)

@view(y[1:end-1]) .= Mul(U , @view(x[2:end]))

mul!(@view(ỹ[1:end-1]) , Ũ , @view(x̃[2:end]))
y[end] = 0
ỹ[end] = 0
x = copy(b)
    y = copy(x)
    _gaussseidel(L,U, b, x, y, 1) - x_bb

x .= b .- y
ỹ .= b .- ỹ
x .= Ldiv(L, x)

x̃ .= L̃ \ ỹ

x - x_bb

x̃-x

(L̃ \ x) - (L \ x)

ỹ - y |> norm
x - x_bb

A*x - b

L̃ == L
Ũ == U

norm(A*x - b)
@time

@time copy(x)
L̃ = sparse(L)
Ũ = sparse(U)

y = x[1:end-1]
z = similar(y)
@time z .= Mul(U,y)
y = similar(x)
@time (@view(y[1:end-1]) .= Mul(U , @view(x[2:end])))
y[end] = 0
@time x .= b .- y
@time x .= Ldiv(L, x)
A = randn(9,40_000);
    @time A*x;

A
S = sparse(A);

400^2

@time L\x

h = 1/n
    @time D² = BandedMatrix(0 => Fill(-2,n), 1 => Fill(1,n-1), -1 => Fill(1,n-1))/h^2
    @time D_xx = BandedBlockBandedMatrix(Kron(D², Eye(n)))
    @time D_yy = BandedBlockBandedMatrix(Kron(Eye(n), D²))
    @time D_xx .+ D_yy
    4

using Profile
dest = Δ;
    @time BlockBandedMatrices.blockbanded_copyto!(dest, D_yy)
    view(dest, Block(1,1)) .= view(dest, Block(1,1)) .+ view(D_yy, Block(1,1))
D_yy
@time copyto!(dest, Broadcast.broadcasted(+, D_xx , D_yy))

using FFTW
A = randn(400, 400);
    @time fft(A);
dest[Block(1,1)]


function compress_mpo!(env, mpstarget, mpo, maxsweeps)

    mps = MPO_to_MPS(mpo)

    compress_var!(mpstarget, mps, env, maxsweeps)

    MPS_to_MPO(mpstarget)

end

"""
function to compress mps A using maxsweeps left and right of the
variational algorithm with Ain as initial guess. Ain must be in left
canonical form (normalisation not necessary)
"""

function compress_var(Ain,A,maxsweeps)

    na = length(Ain)
    FL = Array{Array{Complex{Float64}, 2}, 1}(na)
    FR = Array{Array{Complex{Float64}, 2}, 1}(na)
    A0 = Array{Array{Complex{Float64}, 3}, 1}(na)

    #initializing left environment

    FL[1] = eye(1)
    for jj = 1:(na - 1)
        FL[jj + 1] = update_lenv(Ain[jj], A[jj], FL[jj])
    end
    FR[na] = eye(1)

    for sweep = 1:maxsweeps

        #sweep left

        for jj = na:-1:2
            AR = prod_LR(A[jj], FR[jj])
            A0[jj] = prod_LR(FL[jj], AR)
            A0[jj] = rightorth!(A0[jj])
            FR[jj - 1] = update_renv(A0[jj], AR)
        end
        AR = prod_LR(A[1], FR[1])
        A0[1] = prod_LR(FL[1], AR)

        #sweep right

        for jj = 1:(na - 1)
            AL = prod_LR(FL[jj], A[jj])
            A0[jj] = prod_LR(AL, FR[jj])
            A0[jj] = leftorth!(A0[jj])
            FL[jj + 1] = update_lenv(A0[jj], AL)
        end
        AL = prod_LR(FL[na], A[na])
        A0[na] = prod_LR(AL, FR[na])

    end

    return A0

end

function compress_var!(A0, A, env, maxsweeps)

    na = length(A0)

    #initializing left environment

    env[1][1] = 1.0
    for jj = 1:(na - 1)
        update_lenv_co!(env[jj + 1], A0[jj], A[jj], env[jj])
    end
    env[na + 1][1] = 1.0

    for sweep = 1:maxsweeps

        #sweep left

        for jj = na:-1:2
            AR = prod_LR(A[jj], env[jj + 1])
            prod_LR_co!(A0[jj], env[jj], AR)
            A0[jj] = rightorth!(A0[jj])
            update_renv!(env[jj], A0[jj], AR)
        end
        AR = prod_LR(A[1], env[2])
        prod_LR!(A0[1], env[1], AR)

        #sweep right

        for jj = 1:(na - 1)
            AL = prod_LR_co(env[jj], A[jj])
            prod_LR!(A0[jj], AL, env[jj + 1])
            A0[jj] = leftorth!(A0[jj])
            update_lenv_co!(env[jj + 1], A0[jj], AL)
        end
        AL = prod_LR_co(env[na], A[na])
        prod_LR!(A0[na], AL, env[na + 1])

    end

end
"""
function to compress mps A using maxsweeps left and right of the two
site variational algorithm with Ain as initial guess. Ain must be in left
canonical form (normalisation not necessary)
"""

function compress_var_two_site(Ain,A,maxsweeps,svd_tol,dims,dlim)

    na = length(Ain)
    FL = Array{Array{Complex{Float64},2},1}(na)
    FR = Array{Array{Complex{Float64},2},1}(na)
    A0 = Array{Array{Complex{Float64},3},1}(na)

    #initializing left environment

    FL[1] = eye(1)
    for jj = 1:na-1
        FL[jj+1] = update_lenv(Ain[jj],A[jj],FL[jj])
    end
    FR[na] = eye(1);

    for sweep = 1:maxsweeps

        U,S,Vd = eye(1),eye(1),eye(1)
        #sweep left

        for jj = na:-1:2;

            av = A_prod_LR_two_site(FL[jj-1],A[jj-1],A[jj],FR[jj])
            U,S,Vd = svdecon(reshape(av,dims[jj-1]*size(A[jj-1],2),
                size(A[jj],2)*dims[jj+1]))
            S = S[S.>0]
            stot = sum(S.^2)
            dims[jj] = min(sum(cumsum(S.^2/stot) .< 1-svd_tol)+1,dlim)
            A0[jj] = reshape(Vd[:,1:dims[jj]]',dims[jj],size(A[jj],2),dims[jj+1])
            FR[jj-1] = update_renv(A0[jj],A[jj],FR[jj])

        end
        #A0[1] = reshape(U[:,1:Dims[2]]*diagm(S[1:Dims[2]]),Dims[1],size(A[1],2),Dims[2])

        #sweep right

        for jj = 1:na-1;

            av = A_prod_LR_two_site(FL[jj],A[jj],A[jj+1],FR[jj+1])
            U,S,Vd = svdecon(reshape(av,dims[jj]*size(A[jj],2),
                size(A[jj+1],2)*dims[jj+2]))
            S = S[S.>0]
            stot = sum(S.^2)
            dims[jj+1] = min(sum(cumsum(S.^2/stot) .< 1-svd_tol)+1,dlim)
            A0[jj] = reshape(U[:,1:dims[jj+1]],dims[jj],size(A[jj],2),dims[jj+1])
            FL[jj+1] = update_lenv(A0[jj],A[jj],FL[jj])

        end
        A0[na] = reshape(diagm(S[1:dims[na]])*(Vd[:,1:dims[na]]'),dims[na],size(A[na],2),dims[na+1])

    end

    return A0

end

"""
function to compress mps resulting from the product of H and A using
maxsweeps left and right of the variational algorithm with Ain as
initial guess. Ain must be in left canonical form (normalisation not necessary)
"""

function compress_var_apply_H(Ain, A, H, maxsweeps)

    na = length(Ain)
    FL = Array{Array{Complex{Float64}, 3}, 1}(na)
    FR = Array{Array{Complex{Float64}, 3}, 1}(na)
    A0 = Array{Array{Complex{Float64}, 3}, 1}(na)

    #initializing left environment

    FL[1] = ones(1, 1, 1)
    for jj = 1:(na - 1)
        FL[jj + 1] = update_le(Ain[jj], H[jj], A[jj], FL[jj])
    end
    FR[na] = ones(1, 1, 1)

    for sweep = 1:maxsweeps

        #sweep left

        for jj = na:-1:2;
            AOR = prod_AOR(A[jj], H[jj], FR[jj])
            A0[jj] = prod_LR(FL[jj], AOR)
            A0[jj] = rightorth!(A0[jj])
            FR[jj - 1] = update_re(A0[jj], AOR)
        end

        #sweep right

        for jj = 1:(na - 1)
            AOL = prod_AOL(A[jj], H[jj], FL[jj])
            A0[jj] = prod_LR(AOL, FR[jj])
            A0[jj] = leftorth!(A0[jj])
            FL[jj + 1] = update_le(A0[jj], AOL)
        end

        AOL = prod_AOL(A[na], H[na], FL[na])
        A0[na] = prod_LR(AOL, FR[na])

    end

    return A0

end

function compress_var_apply_H!(A0, envop, Ain, A, H, maxsweeps)

    na = length(Ain)

    #initializing left environment

    envop[1][1] = 1.0
    for jj = 1:(na - 1)
        update_le!(envop[jj + 1], Ain[jj], H[jj], A[jj], envop[jj])
    end
    envop[na + 1][1] = 1.0

    for sweep = 1:maxsweeps

        #sweep left

        for jj = na:-1:2

            AOR = prod_AOR(A[jj], H[jj], envop[jj + 1])
            prod_LR!(A0[jj], envop[jj], AOR)
            A0[jj] = rightorth!(A0[jj])
            update_re!(envop[jj], A0[jj], AOR)

        end

        #sweep right

        for jj = 1:(na - 1)

            AOL = prod_AOL(A[jj], H[jj], envop[jj])
            prod_LR!(A0[jj], AOL, envop[jj + 1])
            A0[jj] = leftorth!(A0[jj])
            update_le!(envop[jj + 1], A0[jj], AOL)

        end

        AOL = prod_AOL(A[na], H[na], envop[na])
        prod_LR!(A0[na], AOL, envop[na + 1])

    end

end

function compress_var_apply_H_right(Ain, A, H, maxsweeps)

    na = length(Ain)
    FR = Array{Array{Complex{Float64}, 3}, 1}(na)
    A0 = Array{Array{Complex{Float64}, 3}, 1}(na)

    #initializing right environment

    FR[na] = ones(1, 1, 1)
    for jj = na:-1:2;
        FR[jj - 1] = update_re(Ain[jj], H[jj], A[jj], FR[jj])
    end
    FL = complex(ones(1, 1, 1))

    #sweep right

    for jj = 1:(na - 1)
        AOL = prod_AOL(A[jj], H[jj], FL)
        A0[jj] = prod_LR(AOL, FR[jj])
        A0[jj] = leftorth!(A0[jj])
        FL = update_le(A0[jj], AOL)
    end

    AOL = prod_AOL(A[na], H[na], FL)
    A0[na] = prod_LR(AOL, FR[na])

    return A0

end

function compress_var_apply_H_left(Ain, A, H)

    na = length(Ain)
    FL = Array{Array{Complex{Float64}, 3}, 1}(na)
    A0 = Array{Array{Complex{Float64}, 3}, 1}(na)

    #initializing left environment

    FL[1] = ones(1, 1, 1)
    for jj = 1:(na - 1)
        FL[jj + 1] = update_le(Ain[jj], H[jj], A[jj], FL[jj])
    end
    FR = complex(ones(1, 1, 1))

    #sweep left

    for jj = na:-1:2;
        AOR = prod_AOR(A[jj], H[jj], FR)
        A0[jj] = prod_LR(FL[jj], AOR)
        A0[jj] = rightorth!(A0[jj])
        FR = update_re(A0[jj], AOR)
    end

    AOR = prod_AOR(A[1], H[1], FR)
    A0[1] = prod_LR(FL[1], AOR)

    return A0

end

function compress_var_apply_H_left!(A0, envop, Ain, A, H)

    na = length(Ain)

    #initializing left environment

    envop[1][1] = 1.0
    for jj = 1:(na - 1)
        update_le!(envop[jj + 1], Ain[jj], H[jj], A[jj], envop[jj])
    end
    envop[na + 1][1] = 1.0

    #sweep left

    for jj = na:-1:2

        AOR = prod_AOR(A[jj], H[jj], envop[jj + 1])
        prod_LR!(A0[jj], envop[jj], AOR)
        A0[jj] = rightorth!(A0[jj])
        update_re!(envop[jj], A0[jj], AOR)

    end

    AOR = prod_AOR(A[1], H[1], envop[2])
    prod_LR!(A0[1], envop[1], AOR)

end

"""
function to compress mps resulting from the product of H and A using
maxsweeps left and right of the variational two site algorithm with Ain as
initial guess. Ain must be in left canonical form (normalisation not necessary)
"""

function compress_var_apply_H_two_site(Ain,A,H,maxsweeps,svd_tol,dims,dlim)

    na = length(Ain)
    FL = Array{Array{Complex{Float64},3},1}(na)
    FR = Array{Array{Complex{Float64},3},1}(na)
    A0 = Array{Array{Complex{Float64},3},1}(na)

    #initializing left environment

    FL[1] = ones(1,1,1)
    for jj = 1:na-2
        FL[jj+1] = update_le(Ain[jj],H[jj],A[jj],FL[jj])
    end
    FR[na] = ones(1,1,1);

    for sweep = 1:maxsweeps

        U,S,Vd = eye(1),eye(1),eye(1)
        #sweep left

        for jj = na:-1:2;

            av = A_prod_LOR_two_site(FL[jj-1],A[jj-1],H[jj-1],A[jj],H[jj],FR[jj])
            U,S,Vd = svdecon(reshape(av,dims[jj-1]*size(A[jj-1],2),
                size(A[jj],2)*dims[jj+1]))
            S = S[S.>0]
            stot = sum(S.^2)
            dims[jj] = min(sum(cumsum(S.^2/stot) .< 1-svd_tol)+1,dlim)
            A0[jj] = reshape(Vd[:,1:dims[jj]]',dims[jj],size(A[jj],2),dims[jj+1])
            FR[jj-1] = update_re(A0[jj],H[jj],A[jj],FR[jj])

        end

        #A0[1] = reshape(U[:,1:Dims[2]]*diagm(S[1:Dims[2]]),Dims[1],size(A[1],2),Dims[2])

        #sweep right

        for jj = 1:na-1;

            av = A_prod_LOR_two_site(FL[jj],A[jj],H[jj],A[jj+1],H[jj+1],FR[jj+1])
            U,S,Vd = svdecon(reshape(av,dims[jj]*size(A[jj],2),
                size(A[jj+1],2)*dims[jj+2]))
            S = S[S.>0]
            stot = sum(S.^2)
            dims[jj+1] = min(sum(cumsum(S.^2/stot) .< 1-svd_tol)+1,dlim)
            A0[jj] = reshape(U[:,1:dims[jj+1]],dims[jj],size(A[jj],2),dims[jj+1])
            FL[jj+1] = update_le(A0[jj],H[jj],A[jj],FL[jj])

        end

        A0[na] = reshape(diagm(S[1:dims[na]])*(Vd[:,1:dims[na]]'),dims[na],size(A[na],2),dims[na+1])

    end

    return A0

end

function compress_var_apply_H_two_site_left(Ain,A,H,maxsweeps,svd_tol,dims,dlim)

    na = length(Ain)
    FL = Array{Array{Complex{Float64},3},1}(na)
    FR = Array{Array{Complex{Float64},3},1}(na)
    A0 = Array{Array{Complex{Float64},3},1}(na)

    #initializing left environment

    FL[1] = ones(1,1,1)
    for jj = 1:na-2
        FL[jj+1] = update_le(Ain[jj],H[jj],A[jj],FL[jj])
    end
    FR[na] = ones(1,1,1);

    U,S,Vd = eye(1),eye(1),eye(1)
    #sweep left

    for jj = na:-1:2;

        av = A_prod_LOR_two_site(FL[jj-1],A[jj-1],H[jj-1],A[jj],H[jj],FR[jj])
        U,S,Vd = svdecon(reshape(av,dims[jj-1]*size(A[jj-1],2),
            size(A[jj],2)*dims[jj+1]))
        S = S[S.>0]
        stot = sum(S.^2)
        dims[jj] = min(sum(cumsum(S.^2/stot) .< 1-svd_tol)+1,dlim)
        A0[jj] = reshape(Vd[:,1:dims[jj]]',dims[jj],size(A[jj],2),dims[jj+1])
        FR[jj-1] = update_re(A0[jj],H[jj],A[jj],FR[jj])

    end

    A0[1] = reshape(U[:,1:dims[2]]*diagm(S[1:dims[2]]),dims[1],size(A[1],2),dims[2])

    return A0

end

function compress_var_apply_H_two_site_right(Ain,A,H,maxsweeps,svd_tol,dims,dlim)

    na = length(Ain)
    FL = Array{Array{Complex{Float64},3},1}(na)
    FR = Array{Array{Complex{Float64},3},1}(na)
    A0 = Array{Array{Complex{Float64},3},1}(na)

    #initializing right environment

    FR[na] = ones(1,1,1);
    for jj = na:-1:2;
        FR[jj-1] = update_re(Ain[jj],H[jj],A[jj],FR[jj])
    end
    FL[1] = ones(1,1,1);

    U,S,Vd = eye(1),eye(1),eye(1)

    #sweep right

    for jj = 1:na-1;

        av = A_prod_LOR_two_site(FL[jj],A[jj],H[jj],A[jj+1],H[jj+1],FR[jj+1])
        U,S,Vd = svdecon(reshape(av,dims[jj]*size(A[jj],2),
            size(A[jj+1],2)*dims[jj+2]))
        S = S[S.>0]
        stot = sum(S.^2)
        dims[jj+1] = min(sum(cumsum(S.^2/stot) .< 1-svd_tol)+1,dlim)
        A0[jj] = reshape(U[:,1:dims[jj+1]],dims[jj],size(A[jj],2),dims[jj+1])
        FL[jj+1] = update_le(A0[jj],H[jj],A[jj],FL[jj])

    end

    A0[na] = reshape(diagm(S[1:dims[na]])*(Vd[:,1:dims[na]]'),dims[na],size(A[na],2),dims[na+1])

    return A0

end

"""
function to compress linear combination of mps stored in array A, which
have weights given by coef, using a
variational algorithm with Ain as initial guess. Ain must be in left
canonical form (normalisation not necessary)
"""

function compress_sum_var(Ain, coef, A, maxsweeps)

    na = length(Ain)
    numv = length(coef)
    FL = Array{Array{Array{Complex{Float64}, 2}, 1}, 1}(numv)
    FR = Array{Array{Array{Complex{Float64}, 2}, 1}, 1}(numv)
    A0 = Array{Array{Complex{Float64}, 3}, 1}(na)
    AR = Array{Array{Complex{Float64}, 3}, 1}(numv)

    #initializing left environment

    for ll = 1:numv
        #A[ll][na] = coef[ll]*Av[ll][na]
        FL[ll] = Array{Array{Complex{Float64}, 2}, 1}(na)
        FL[ll][1] = eye(1)
        for jj = 1:(na - 1)
            FL[ll][jj + 1] = update_lenv(Ain[jj], A[ll][jj], FL[ll][jj])
        end
        FR[ll] = Array{Array{Complex{Float64}, 2}, 1}(na)
        FR[ll][na] = coef[ll]*eye(1)
    end

    for sweep = 1:maxsweeps

        #sweep left

        for jj = na:-1:2

            AR[1] = prod_LR(A[1][jj], FR[1][jj])
            A0[jj] = prod_LR(FL[1][jj], AR[1])
            for ll = 2:numv
                AR[ll] = prod_LR(A[ll][jj], FR[ll][jj])
                A0[jj] += prod_LR(FL[ll][jj], AR[ll])
            end
            A0[jj] = rightorth!(A0[jj])
            for ll = 1:numv
                FR[ll][jj - 1] = update_renv(A0[jj], AR[ll])
            end

        end

        #sweep right

        for jj = 1:(na - 1)

            AR[1] = prod_LR(FL[1][jj], A[1][jj])
            A0[jj] = prod_LR(AR[1], FR[1][jj])
            for ll = 2:numv
                AR[ll] = prod_LR(FL[ll][jj], A[ll][jj])
                A0[jj] += prod_LR(AR[ll], FR[ll][jj])
            end
            A0[jj] = leftorth!(A0[jj])
            for ll = 1:numv
                FL[ll][jj + 1] = update_lenv(A0[jj], AR[ll])
            end

        end

    end

    AR[1] = prod_LR(FL[1][na], A[1][na])
    A0[na] = prod_LR(AR[1], FR[1][na])
    for ll = 2:numv
        AR[ll] = prod_LR(FL[ll][na], A[ll][na])
        A0[na] += prod_LR(AR[ll], FR[ll][na])
    end

    return A0

end

"""
function to compress linear combination of mps stored array A, which
have weights given by coef, using the two site
variational algorithm with Ain as initial guess. Ain must be in left
canonical form (normalisation not necessary)
"""

function compress_sum_var_two_site(Ain,coef,A,maxsweeps,svd_tol,dims,dlim)

    na = length(Ain)
    numv = length(coef)
    FL = Array{Array{Array{Complex{Float64},2},1},1}(numv)
    FR = Array{Array{Array{Complex{Float64},2},1},1}(numv)
    A0 = Array{Array{Complex{Float64},3},1}(na)

    #initializing left environment

    for ll = 1:numv

        #A[ll][na] = coef[ll]*Av[ll][na]
        FL[ll] = Array{Array{Complex{Float64},2},1}(na)
        FL[ll][1] = eye(1)
        for jj = 1:na-2
            FL[ll][jj+1] = update_lenv(Ain[jj],A[ll][jj],FL[ll][jj])
        end
        FR[ll] = Array{Array{Complex{Float64},2},1}(na)
        FR[ll][na] = coef[ll]*eye(1);

    end

    for sweep = 1:maxsweeps

        U,S,Vd = eye(1),eye(1),eye(1)
        #sweep left

        for jj = na:-1:2

            av = A_prod_LR_two_site(FL[1][jj-1],A[1][jj-1],A[1][jj],FR[1][jj])
            for ll = 2:numv
                av += A_prod_LR_two_site(FL[ll][jj-1],A[ll][jj-1],A[ll][jj],FR[ll][jj])
            end
            U,S,Vd = svdecon(reshape(av,dims[jj-1]*size(A[1][jj-1],2),
                size(A[1][jj],2)*dims[jj+1]))
            S = S[S.>0]
            stot = sum(S.^2)
            dims[jj] = min(sum(cumsum(S.^2/stot) .< 1-svd_tol)+1,dlim)
            A0[jj] = reshape(Vd[:,1:dims[jj]]',dims[jj],size(A[1][jj],2),dims[jj+1])
            for ll = 1:numv
                FR[ll][jj-1] = update_renv(A0[jj],A[ll][jj],FR[ll][jj])
            end

        end

        #A0[1] = reshape(U[:,1:Dims[2]]*diagm(S[1:Dims[2]]),Dims[1],size(A[1][1],2),Dims[2])

        #sweep right

        for jj = 1:na-1;

            av = A_prod_LR_two_site(FL[1][jj],A[1][jj],A[1][jj+1],FR[1][jj+1])
            for ll = 2:numv
                av += A_prod_LR_two_site(FL[ll][jj],A[ll][jj],A[ll][jj+1],FR[ll][jj+1])
            end
            U,S,Vd = svdecon(reshape(av,dims[jj]*size(A[1][jj],2),
                size(A[1][jj+1],2)*dims[jj+2]))
            S = S[S.>0]
            stot = sum(S.^2)
            dims[jj+1] = min(sum(cumsum(S.^2/stot) .< 1-svd_tol)+1,dlim)
            A0[jj] = reshape(U[:,1:dims[jj+1]],dims[jj],size(A[1][jj],2),dims[jj+1])
            for ll = 1:numv
                FL[ll][jj+1] = update_lenv(A0[jj],A[ll][jj],FL[ll][jj])
            end
        end

        A0[na] = reshape(diagm(S[1:dims[na]])*(Vd[:,1:dims[na]]'),dims[na],size(A[1][na],2),dims[na+1])

    end

    return A0

end

"""
function to compress linear combination of mps stored in cell array A, which
have weights given by coef and where H is applied to the last MPS.
The compression uses the variational algorithm with A0 as initial guess.
A0 must be in left canonical form (normalisation not necessary)
"""

function compress_sum_var_apply_H(Ain, coef, A, H, maxsweeps)

    na = length(Ain)
    numv = length(coef)
    FL = Array{Array{Array{Complex{Float64}, 2}, 1}, 1}(numv - 1)
    FR = Array{Array{Array{Complex{Float64}, 2}, 1}, 1}(numv - 1)
    FLH = Array{Array{Complex{Float64}, 3}, 1}(na)
    FRH = Array{Array{Complex{Float64} ,3}, 1}(na)
    A0 = Array{Array{Complex{Float64}, 3}, 1}(na)
    AR = Array{Array{Complex{Float64}, 3}, 1}(numv - 1)

    #initializing left environment

    for ll = 1:(numv - 1)
        #A[ll][na] = coef[ll]*Av[ll][na]
        FL[ll] = Array{Array{Complex{Float64}, 2}, 1}(na)
        FL[ll][1] = eye(1)
        for jj = 1:(na - 1)
            FL[ll][jj + 1] = update_lenv(Ain[jj], A[ll][jj], FL[ll][jj])
        end
        FR[ll] = Array{Array{Complex{Float64}, 2}, 1}(na)
        FR[ll][na] = coef[ll]*eye(1)
    end

    FLH[1] = ones(1, 1, 1)
    for jj = 1:(na - 1)
        FLH[jj + 1] = update_le(Ain[jj], H[jj], A[numv][jj], FLH[jj])
    end
    FRH[na] = coef[numv]*ones(1, 1, 1)

    for sweep = 1:maxsweeps

        #sweep left

        for jj = na:-1:2;

            AR[1] = prod_LR(A[1][jj], FR[1][jj])
            A0[jj] = prod_LR(FL[1][jj], AR[1])
            for ll = 2:(numv - 1)
                AR[ll] = prod_LR(A[ll][jj], FR[ll][jj])
                A0[jj] += prod_LR(FL[ll][jj], AR[ll])
            end
            AOR = prod_AOR(A[numv][jj], H[jj], FRH[jj])
            A0[jj] += prod_LR(FLH[jj], AOR)
            A0[jj] = rightorth!(A0[jj])
            for ll = 1:(numv - 1)
                FR[ll][jj - 1] = update_renv(A0[jj], AR[ll])
            end
            FRH[jj - 1] = update_re(A0[jj], AOR)
        end

        #sweep right

        for jj = 1:(na - 1)

            AR[1] = prod_LR(FL[1][jj], A[1][jj])
            A0[jj] = prod_LR(AR[1], FR[1][jj])
            for ll = 2:(numv - 1)
                AR[ll] = prod_LR(FL[ll][jj], A[ll][jj])
                A0[jj] += prod_LR(AR[ll], FR[ll][jj])
            end
            AOL = prod_AOL(A[numv][jj], H[jj], FLH[jj])
            A0[jj] += prod_LR(AOL, FRH[jj])
            A0[jj] = leftorth!(A0[jj])
            for ll = 1:(numv - 1)
                FL[ll][jj + 1] = update_lenv(A0[jj], AR[ll])
            end
            FLH[jj + 1] = update_le(A0[jj], AOL)
        end

    end

    AR[1] = prod_LR(FL[1][na], A[1][na])
    A0[na] = prod_LR(AR[1], FR[1][na])
    for ll = 2:numv-1
        AR[ll] = prod_LR(FL[ll][na], A[ll][na])
        A0[na] += prod_LR(AR[ll], FR[ll][na])
    end
    AOL = prod_AOL(A[numv][na], H[na], FLH[na])
    A0[na] += prod_LR(AOL, FRH[na])

    return A0

end

function compress_sum_var_apply_H_left(Ain, coef, A, H)

    na = length(Ain)
    numv = length(coef)
    FL = Array{Array{Array{Complex{Float64}, 2}, 1}, 1}(numv - 1)
    FR = Array{Array{Complex{Float64}, 2}, 1}(numv - 1)
    FLH = Array{Array{Complex{Float64}, 3}, 1}(na)
    A0 = Array{Array{Complex{Float64}, 3}, 1}(na)
    AR = Array{Array{Complex{Float64}, 3}, 1}(numv - 1)

    #initializing left environment

    for ll = 1:(numv - 1)
        #A[ll][na] = coef[ll]*Av[ll][na]
        FL[ll] = Array{Array{Complex{Float64}, 2}, 1}(na)
        FL[ll][1] = eye(1)
        for jj = 1:(na - 1)
            FL[ll][jj + 1] = update_lenv(Ain[jj], A[ll][jj], FL[ll][jj])
        end
        FR[ll] = coef[ll]*eye(1)

    end

    FLH[1] = ones(1, 1, 1)
    for jj = 1:(na - 1)
        FLH[jj + 1] = update_le(Ain[jj], H[jj], A[numv][jj], FLH[jj])
    end

    FRH = coef[numv]*complex(ones(1, 1, 1))

    #sweep left

    for jj = na:-1:2;

        AR[1] = prod_LR(A[1][jj], FR[1])
        A0[jj] = prod_LR(FL[1][jj], AR[1])
        for ll = 2:(numv - 1)
            AR[ll] = prod_LR(A[ll][jj], FR[ll])
            A0[jj] += prod_LR(FL[ll][jj], AR[ll])
        end
        AOR = prod_AOR(A[numv][jj], H[jj], FRH)
        A0[jj] += prod_LR(FLH[jj], AOR)
        A0[jj] = rightorth!(A0[jj])
        for ll = 1:(numv - 1)
            FR[ll] = update_renv(A0[jj], AR[ll])
        end
        FRH = update_re(A0[jj], AOR)
    end

    AR[1] = prod_LR(A[1][1], FR[1])
    A0[1] = prod_LR(FL[1][1], AR[1])
    for ll = 2:numv-1
        AR[ll] = prod_LR(A[ll][1], FR[ll])
        A0[1] += prod_LR(FL[ll][1], AR[ll])
    end
    AOR = prod_AOR(A[numv][1], H[1], FRH)
    A0[1] += prod_LR(FLH[1], AOR)

    return A0

end

function compress_sum_var_apply_H!(A0, envop, coef, A, H, maxsweeps)

    na = length(A0)
    numv = length(coef)
    env = Array{Array{Array{Complex{Float64}, 2}, 1}, 1}(numv - 1)
    AR = Array{Array{Complex{Float64}, 3}, 1}(numv - 1)

    #initializing left environment

    for ll = 1:(numv - 1)
        #A[ll][na] = coef[ll]*Av[ll][na]
        env[ll] = Array{Array{Complex{Float64}, 2}, 1}(na + 1)
        env[ll][1][1] = 1.0
        for jj = 1:(na - 1)
            env[ll][jj + 1] = update_lenv(A0[jj], A[ll][jj], env[ll][jj])
        end
        env[ll][na + 1][1] = coef[ll]
    end

    envop[1][1] = 1.0
    for jj = 1:(na - 1)

        update_le!(envop[jj + 1], A0[jj], H[jj], A[numv][jj], envop[jj])

    end
    envop[na + 1][1] = coef[numv]

    for sweep = 1:maxsweeps

        #sweep left

        for jj = na:-1:2;

            AR[1] = prod_LR(A[1][jj], env[1][jj + 1])
            prod_LR!(A0[jj], env[1][jj], AR[1])
            for ll = 2:(numv - 1)
                AR[ll] = prod_LR(A[ll][jj], env[ll][jj + 1])
                A0[jj] += prod_LR(env[ll][jj], AR[ll])
            end
            AOR = prod_AOR(A[numv][jj], H[jj], envop[jj + 1])
            A0[jj] += prod_LR(envop[jj], AOR)
            A0[jj] = rightorth!(A0[jj])
            for ll = 1:(numv - 1)
                update_renv!(env[ll][jj], A0[jj], AR[ll])
            end
            update_re!(envop[jj], A0[jj], AOR)

        end

        #sweep right

        for jj = 1:(na - 1)

            AR[1] = prod_LR(env[1][jj], A[1][jj])
            prod_LR!(A0[jj], AR[1], env[1][jj + 1])
            for ll = 2:(numv - 1)
                AR[ll] = prod_LR(env[ll][jj], A[ll][jj])
                A0[jj] += prod_LR(AR[ll], env[ll][jj + 1])
            end
            AOL = prod_AOL(A[numv][jj], H[jj], envop[jj])
            A0[jj] += prod_LR(AOL, envop[jj + 1])
            A0[jj] = leftorth!(A0[jj])
            for ll = 1:(numv - 1)
                update_lenv!(env[ll][jj + 1], A0[jj], AR[ll])
            end
            update_le!(envop[jj + 1], A0[jj], AOL)

        end

    end

    AR[1] = prod_LR(env[1][na], A[1][na])
    prod_LR!(A0[na], AR[1], env[1][na + 1])
    for ll = 2:numv-1
        AR[ll] = prod_LR(env[ll][na], A[ll][na])
        A0[na] += prod_LR(AR[ll], env[ll][na + 1])
    end
    AOL = prod_AOL(A[numv][na], H[na], envop[na])
    A0[na] += prod_LR(AOL, envop[na + 1])

    nothing

end

function compress_sum_var_apply_H_left!(A0, env, envop, Ain, coef, A, H)

    na = length(Ain)
    numv = length(coef)
    AR = similar(A0, numv - 1)

    #initializing left environment

    for ll = 1:(numv - 1)

        env[ll][1][1] = 1.0
        for jj = 1:(na - 1)
            update_lenv!(env[ll][jj + 1], Ain[jj], A[ll][jj], env[ll][jj])
        end
        env[ll][na + 1][1] = coef[ll]

    end

    envop[1][1] = 1.0
    for jj = 1:(na - 1)

        update_le!(envop[jj + 1], Ain[jj], H[jj], A[numv][jj], envop[jj])

    end
    envop[na + 1][1] = coef[numv]

    #sweep left

    for jj = na:-1:2;

        AR[1] = prod_LR(A[1][jj], env[1][jj + 1])
        prod_LR!(A0[jj], env[1][jj], AR[1])
        for ll = 2:(numv - 1)
            AR[ll] = prod_LR(A[ll][jj], env[ll][jj + 1])
            A0[jj] += prod_LR(env[ll][jj], AR[ll])
        end
        AOR = prod_AOR(A[numv][jj], H[jj], envop[jj + 1])
        A0[jj] += prod_LR(envop[jj], AOR)
        A0[jj] = rightorth!(A0[jj])
        for ll = 1:(numv - 1)
            update_renv!(env[ll][jj], A0[jj], AR[ll])
        end
        update_re!(envop[jj], A0[jj], AOR)
    end

    AR[1] = prod_LR(A[1][1], env[1][2])
    prod_LR!(A0[1], env[1][1], AR[1])
    for ll = 2:numv-1
        AR[ll] = prod_LR(A[ll][1], env[ll][2])
        A0[1] += prod_LR(env[ll][1], AR[ll])
    end
    AOR = prod_AOR(A[numv][1], H[1], envop[2])
    A0[1] += prod_LR(envop[1], AOR)
    nothing

end

function compress_sum_var_apply_H_right(Ain, coef, A, H)

    na = length(Ain)
    numv = length(coef)
    FL = Array{Array{Complex{Float64}, 2}, 1}(numv - 1)
    FR = Array{Array{Array{Complex{Float64}, 2}, 1}, 1}(numv - 1)
    FRH = Array{Array{Complex{Float64}, 3}, 1}(na)
    A0 = Array{Array{Complex{Float64}, 3}, 1}(na)
    AL = Array{Array{Complex{Float64}, 3}, 1}(numv - 1)

    #initializing right environment

    for ll = 1:(numv - 1)
        FR[ll] = Array{Array{Complex{Float64}, 2}, 1}(na)
        FR[ll][na] = eye(1)
        for jj = na:-1:2
            FR[ll][jj - 1] = update_renv(Ain[jj], A[ll][jj], FR[ll][jj])
        end
        FL[ll] = coef[ll]*eye(1)
    end

    FRH[na] = ones(1, 1, 1)
    for jj = na:-1:2
        FRH[jj - 1] = update_re(Ain[jj], H[jj], A[numv][jj], FRH[jj])
    end

    FLH = coef[numv]*complex(ones(1, 1, 1))

    #sweep right

    for jj = 1:(na - 1)

        AL[1] = prod_LR(FL[1], A[1][jj])
        A0[jj] = prod_LR(AL[1], FR[1][jj])
        for ll = 2:(numv - 1)
            AL[ll] = prod_LR(FL[ll], A[ll][jj])
            A0[jj] += prod_LR(AL[ll], FR[ll][jj])
        end
        AOL = prod_AOL(A[numv][jj], H[jj], FLH)
        A0[jj] += prod_LR(AOL, FRH[jj])
        A0[jj] = leftorth!(A0[jj])
        for ll = 1:(numv - 1)
            FL[ll] = update_lenv(A0[jj], AL[ll])
        end
        FLH = update_le(A0[jj], AOL)

    end

    AL[1] = prod_LR(FL[1], A[1][na])
    A0[na] = prod_LR(AL[1], FR[1][na])
    for ll = 2:(numv - 1)
        AL[ll] = prod_LR(FL[ll], A[ll][na])
        A0[na] += prod_LR(AL[ll], FR[ll][na])
    end
    AOL = prod_AOL(A[numv][na], H[na], FLH)
    A0[na] += prod_LR(AOL, FRH[na])

    return A0

end

function compress_sum_var_apply_H_right!(A0, env, envop, Ain, coef, A, H)

    na = length(Ain)
    numv = length(coef)
    AL = similar(A0, numv - 1)

    #initializing right environment

    for ll = 1:(numv - 1)

        env[ll][na + 1][1] = 1.0
        for jj = na:-1:2
            update_renv!(env[ll][jj], Ain[jj], A[ll][jj], env[ll][jj + 1])
        end
        env[ll][1][1] = coef[ll]

    end

    envop[na + 1][1] = 1.0
    for jj = na:-1:2

        update_re!(envop[jj], Ain[jj], H[jj], A[numv][jj], envop[jj + 1])

    end
    envop[1][1] = coef[numv]

    #sweep right

    for jj = 1:(na - 1)

        AL[1] = prod_LR(env[1][jj], A[1][jj])
        prod_LR!(A0[jj], AL[1], env[1][jj + 1])
        for ll = 2:(numv - 1)
            AL[ll] = prod_LR(env[ll][jj], A[ll][jj])
            A0[jj] += prod_LR(AL[ll], env[ll][jj + 1])
        end
        AOL = prod_AOL(A[numv][jj], H[jj], envop[jj])
        A0[jj] += prod_LR(AOL, envop[jj + 1])
        A0[jj] = leftorth!(A0[jj])
        for ll = 1:(numv - 1)
            update_lenv!(env[ll][jj + 1], A0[jj], AL[ll])
        end
        update_le!(envop[jj + 1], A0[jj], AOL)

    end

    AL[1] = prod_LR(env[1][na], A[1][na])
    prod_LR!(A0[na], AL[1], env[1][na + 1])
    for ll = 2:(numv - 1)
        AL[ll] = prod_LR(env[ll][na], A[ll][na])
        A0[na] += prod_LR(AL[ll], env[ll][na + 1])
    end
    AOL = prod_AOL(A[numv][na], H[na], envop[na])
    A0[na] += prod_LR(AOL, envop[na + 1])
    nothing

end

"""
function to compress linear combination of mps stored in array A, which
have weights given by coef and where H is applied to the last MPS.
The compression uses the two site variational algorithm with A0 as initial guess.
A0 must be in left canonical form (normalisation not necessary)
"""

function compress_sum_var_apply_H_two_site(Ain,coef,A,H,maxsweeps,svd_tol,dims,dlim)

    na = length(Ain)
    numv = length(coef)
    FL = Array{Array{Array{Complex{Float64},2},1},1}(numv-1)
    FR = Array{Array{Array{Complex{Float64},2},1},1}(numv-1)
    FLH = Array{Array{Complex{Float64},3},1}(na)
    FRH = Array{Array{Complex{Float64},3},1}(na)
    A0 = Array{Array{Complex{Float64},3},1}(na)

    #initializing left environment

    for ll = 1:numv-1
        #A[ll][na] = coef[ll]*Av[ll][na]
        FL[ll] = Array{Array{Complex{Float64},2},1}(na)
        FL[ll][1] = eye(1)
        for jj = 1:na-2
            FL[ll][jj+1] = update_lenv(Ain[jj],A[ll][jj],FL[ll][jj])
        end
        FR[ll] = Array{Array{Complex{Float64},2},1}(na)
        FR[ll][na] = coef[ll]*eye(1);
    end

    FLH[1] = ones(1,1,1)
    for jj = 1:na-1
        FLH[jj+1] = update_le(Ain[jj],H[jj],A[numv][jj],FLH[jj])
    end
    FRH[na] = coef[numv]*ones(1,1,1)

    for sweep = 1:maxsweeps

        U,S,Vd = eye(1),eye(1),eye(1)
        #sweep left

        for jj = na:-1:2;

            av = A_prod_LR_two_site(FL[1][jj-1],A[1][jj-1],A[1][jj],FR[1][jj])
            for ll = 2:numv-1
                av += A_prod_LR_two_site(FL[ll][jj-1],A[ll][jj-1],A[ll][jj],FR[ll][jj])
            end
            av += A_prod_LOR_two_site(FLH[jj-1],A[numv][jj-1],H[jj-1],A[numv][jj],H[jj],FRH[jj])
            U,S,Vd = svdecon(reshape(av,dims[jj-1]*size(A[1][jj-1],2),
                size(A[1][jj],2)*dims[jj+1]))
            S = S[S.>0]
            stot = sum(S.^2)
            dims[jj] = min(sum(cumsum(S.^2/stot) .< 1-svd_tol)+1,dlim)
            A0[jj] = reshape(Vd[:,1:dims[jj]]',dims[jj],size(A[1][jj],2),dims[jj+1])
            for ll = 1:numv-1
                FR[ll][jj-1] = update_renv(A0[jj],A[ll][jj],FR[ll][jj])
            end
            FRH[jj-1] = update_re(A0[jj],H[jj],A[numv][jj],FRH[jj])

        end

        #A0[1] = reshape(U[:,1:Dims[2]]*diagm(S[1:Dims[2]]),Dims[1],size(A[1][1],2),Dims[2])

        #sweep right

        for jj = 1:na-1;

            av = A_prod_LR_two_site(FL[1][jj],A[1][jj],A[1][jj+1],FR[1][jj+1])
            for ll = 2:numv-1
                av += A_prod_LR_two_site(FL[ll][jj],A[ll][jj],A[ll][jj+1],FR[ll][jj+1])
            end
            av += A_prod_LOR_two_site(FLH[jj],A[numv][jj],H[jj],A[numv][jj+1],H[jj+1],FRH[jj+1])
            U,S,Vd = svdecon(reshape(av,dims[jj]*size(A[1][jj],2),
                size(A[1][jj+1],2)*dims[jj+2]))
            S = S[S.>0]
            stot = sum(S.^2)
            dims[jj+1] = min(sum(cumsum(S.^2/stot) .< 1-svd_tol)+1,dlim)
            A0[jj] = reshape(U[:,1:dims[jj+1]],dims[jj],size(A[1][jj],2),dims[jj+1])
            for ll = 1:numv-1
                FL[ll][jj+1] = update_lenv(A0[jj],A[ll][jj],FL[ll][jj])
            end
            FLH[jj+1] = update_le(A0[jj],H[jj],A[numv][jj],FLH[jj])

        end

        A0[na] = reshape(diagm(S[1:dims[na]])*(Vd[:,1:dims[na]]'),dims[na],size(A[1][na],2),dims[na+1])

    end

    return A0

end

function compress_sum_var_apply_H_two_site_left(Ain,coef,A,H,maxsweeps,svd_tol,dims,dlim)

    na = length(Ain)
    numv = length(coef)
    FL = Array{Array{Array{Complex{Float64},2},1},1}(numv-1)
    FR = Array{Array{Array{Complex{Float64},2},1},1}(numv-1)
    FLH = Array{Array{Complex{Float64},3},1}(na)
    FRH = Array{Array{Complex{Float64},3},1}(na)
    A0 = Array{Array{Complex{Float64},3},1}(na)

    #initializing left environment

    for ll = 1:numv-1
        #A[ll][na] = coef[ll]*Av[ll][na]
        FL[ll] = Array{Array{Complex{Float64},2},1}(na)
        FL[ll][1] = eye(1)
        for jj = 1:na-2
            FL[ll][jj+1] = update_lenv(Ain[jj],A[ll][jj],FL[ll][jj])
        end
        FR[ll] = Array{Array{Complex{Float64},2},1}(na)
        FR[ll][na] = coef[ll]*eye(1);
    end

    FLH[1] = ones(1,1,1)
    for jj = 1:na-2
        FLH[jj+1] = update_le(Ain[jj],H[jj],A[numv][jj],FLH[jj])
    end
    FRH[na] = coef[numv]*ones(1,1,1)

    U,S,Vd = eye(1),eye(1),eye(1)
    #sweep left

    for jj = na:-1:2;

        av = A_prod_LR_two_site(FL[1][jj-1],A[1][jj-1],A[1][jj],FR[1][jj])
        for ll = 2:numv-1
            av += A_prod_LR_two_site(FL[ll][jj-1],A[ll][jj-1],A[ll][jj],FR[ll][jj])
        end
        av += A_prod_LOR_two_site(FLH[jj-1],A[numv][jj-1],H[jj-1],A[numv][jj],H[jj],FRH[jj])
        U,S,Vd = svdecon(reshape(av,dims[jj-1]*size(A[1][jj-1],2),
            size(A[1][jj],2)*dims[jj+1]))
        S = S[S.>0]
        stot = sum(S.^2)
        dims[jj] = min(sum(cumsum(S.^2/stot) .< 1-svd_tol)+1,dlim)
        A0[jj] = reshape(Vd[:,1:dims[jj]]',dims[jj],size(A[1][jj],2),dims[jj+1])
        for ll = 1:numv-1
            FR[ll][jj-1] = update_renv(A0[jj],A[ll][jj],FR[ll][jj])
        end
        FRH[jj-1] = update_re(A0[jj],H[jj],A[numv][jj],FRH[jj])

    end

    A0[1] = reshape(U[:,1:dims[2]]*diagm(S[1:dims[2]]),dims[1],size(A[1][1],2),dims[2])

    return A0

end

function compress_sum_var_apply_H_two_site_right(Ain,coef,A,H,maxsweeps,svd_tol,dims,dlim)

    na = length(Ain)
    numv = length(coef)
    FL = Array{Array{Array{Complex{Float64},2},1},1}(numv-1)
    FR = Array{Array{Array{Complex{Float64},2},1},1}(numv-1)
    FLH = Array{Array{Complex{Float64},3},1}(na)
    FRH = Array{Array{Complex{Float64},3},1}(na)
    A0 = Array{Array{Complex{Float64},3},1}(na)

    #initializing right environment

    for ll = 1:numv-1
        FR[ll] = Array{Array{Complex{Float64},2},1}(na)
        FR[ll][na] = eye(1)
        for jj = na:-1:2
            FR[ll][jj-1] = update_renv(Ain[jj],A[ll][jj],FR[ll][jj])
        end
        FL[ll] = Array{Array{Complex{Float64},2},1}(na)
        FL[ll][1] = coef[ll]*eye(1)
    end

    FLH[1] = coef[numv]*ones(1,1,1)
    FRH[na] = ones(1,1,1)
    for jj = na:-1:2;
        FRH[jj-1] = update_re(Ain[jj],H[jj],A[numv][jj],FRH[jj])
    end

    U,S,Vd = eye(1),eye(1),eye(1)
    #sweep right

    for jj = 1:na-1;

        av = A_prod_LR_two_site(FL[1][jj],A[1][jj],A[1][jj+1],FR[1][jj+1])
        for ll = 2:numv-1
            av += A_prod_LR_two_site(FL[ll][jj],A[ll][jj],A[ll][jj+1],FR[ll][jj+1])
        end
        av += A_prod_LOR_two_site(FLH[jj],A[numv][jj],H[jj],A[numv][jj+1],H[jj+1],FRH[jj+1])
        U,S,Vd = svdecon(reshape(av,dims[jj]*size(A[1][jj],2),
            size(A[1][jj+1],2)*dims[jj+2]))
        S = S[S.>0]
        stot = sum(S.^2)
        dims[jj+1] = min(sum(cumsum(S.^2/stot) .< 1-svd_tol)+1,dlim)
        A0[jj] = reshape(U[:,1:dims[jj+1]],dims[jj],size(A[1][jj],2),dims[jj+1])
        for ll = 1:numv-1
            FL[ll][jj+1] = update_lenv(A0[jj],A[ll][jj],FL[ll][jj])
        end
        FLH[jj+1] = update_le(A0[jj],H[jj],A[numv][jj],FLH[jj])

    end

    A0[na] = reshape(diagm(S[1:dims[na]])*(Vd[:,1:dims[na]]'),dims[na],size(A[1][na],2),dims[na+1])

    return A0

end

"""
Compresses a left-canonical MPS to a maximum bond dimension D_max using a
relative tolerance tol in the trimming, giving a right-canonical MPS.
"""
function compressionSVD!(mps, dmax, tol = 1e-15)

    n = length(mps)
    lambda = Array{Array}(n - 1)
    err_trim = Array{Float64}(n)
    err_glob = Float64
    ent_ent = Array{Float64}(n)
    dims = Array{Int64}(n - 1)
    dnew = 1
    r = 1

    for i = n:-1:1

        dl, d, dr = size(mps[i])
        mps_site = reshape(mps[i], (dl*d, dr))*r
        #u, s, v = svdecon(reshape(mps_site, (dl, dnew*d)))
        u, s, v = svd(reshape(mps_site, (dl, dnew*d)))
        s = s[s .> 0]
        stot = sum(abs2.(s))
        b = dnew
        dnew = minimum([sum(cumsum(abs2.(s)/stot) .< (1 - tol)) + 1, dmax])
        err_trim[i] = sum(abs2.(s[(dnew + 1):end]))/stot
        ent_ent[i] = -2*sum(abs2.(s).*log.(s))
        mps[i] = reshape(v[:, 1:dnew]', (dnew, d, b))
        r = u[:, 1:dnew]*diagm(s[1:dnew])

    end

    err_glob = 1 - abs2(r[1])

    return err_glob, err_trim, ent_ent

end


@testset "SBitstream" begin
    x = SBitstream(0.5)
    y = SBitstream(0.1)
    w = SBitstream(-0.1)
    X = SBitstream(0.5 * rand(2, 2))
    Y = SBitstream(-0.5 * rand(2, 2))
    v = SBitstream([0.5 * rand(), -0.5 * rand()])

    @testset "FP Values" begin
        @test x.value == 0.5 && y.value == 0.1
        @test float(x) == x.value
        @testset for op in (+, -, *)
            @test op(x, y).value == op(x.value, y.value)
            @test op(x, w).value == op(x.value, w.value)
        end
        @test (y / x).value == y.value / x.value
        @test_logs (:warn, "SBitstream can only be ∈ [-1, 1] (saturation occurring).") (x / y).value == 1
        @test (x ÷ 2).value == x.value / 2
        @test (w ÷ 2).value == w.value / 2
        @test_throws ErrorException x ÷ 0.2
        @test sqrt(x).value == sqrt(x.value)
        @test decorrelate(x).value == x.value
        @test float.(X * Y) == float.(X) * float.(Y)
        @test float.(X * v) == float.(X) * float.(v)
        @test norm(v).value == norm(float.(v))
    end

    @testset "Sample Generation" begin
        xbits = generate(x, 100000)
        @test mean(map(xbit -> pos(xbit) - neg(xbit), xbits)) ≈ x.value atol = 0.01
    end

    @testset "Bit-Level Ops" begin
        T = 30000
        @testset for op in (+, -, *, /, ÷)
            if op == ÷
                @test begin
                    z = 0
                    for t in 1:T
                        bit = pop!(op(x, 2))
                        z += pos(bit) - neg(bit)
                    end

                    z / T
                end ≈ op(x, 2).value atol = 0.01
                @test begin
                    z = 0
                    for t in 1:T
                        bit = pop!(op(w, 2))
                        z += pos(bit) - neg(bit)
                    end

                    z / T
                end ≈ op(w, 2).value atol = 0.01
            elseif op == /
                @test begin
                    z = 0
                    for t in 1:T
                        bit = pop!(op(y, x))
                        z += pos(bit) - neg(bit)
                    end

                    z / T
                end ≈ op(y, x).value atol = 0.01
            else
                @test begin
                    z = 0
                    for t in 1:T
                        bit = pop!(op(x, y))
                        z += pos(bit) - neg(bit)
                    end

                    z / T
                end ≈ op(x, y).value atol = 0.01
                @test begin
                    z = 0
                    for t in 1:T
                        bit = pop!(op(y, x))
                        z += pos(bit) - neg(bit)
                    end

                    z / T
                end ≈ op(y, x).value atol = 0.01
                @test begin
                    z = 0
                    for t in 1:T
                        bit = pop!(op(x, w))
                        z += pos(bit) - neg(bit)
                    end

                    z / T
                end ≈ op(x, w).value atol = 0.01
            end
        end
        @testset for op in (sqrt, decorrelate)
            @test begin
                z = 0
                for t in 1:T
                    bit = pop!(op(x))
                    z += pos(bit) - neg(bit)
                end

                z / T
            end ≈ op(x).value atol = 0.01
        end
        @testset "op = * (matrix, matrix)" begin
            @test begin
                Z = zeros(2, 2)
                for t in 1:T
                    bit = pop!.(X * Y)
                    Z .+= pos.(bit) .- neg.(bit)
                end

                Z ./ T
            end ≈ float.(X * Y) atol = 0.05
            @test begin
                Z = zeros(2, 2)
                for t in 1:T
                    bit = pop!.(Y * X)
                    Z .+= pos.(bit) .- neg.(bit)
                end

                Z ./ T
            end ≈ float.(Y * X) atol = 0.05
        end
        @testset "op = * (matrix, vector)" begin
            @test begin
                T = 50000
                Z = zeros(2)
                for t in 1:T
                    bit = pop!.(X * v)
                    Z .+= pos.(bit) .- neg.(bit)
                end

                Z ./ T
            end ≈ float.(X * v) atol = 0.05
            @test_broken begin
                T = 100000
                Z = zeros(1, 2)
                for t in 1:T
                    bit = pop!.(transpose(v) * X)
                    Z .+= pos.(bit) .- neg.(bit)
                end

                Z ./ T
            end ≈ float.(transpose(v) * X) atol = 0.05
        end
        @testset "op = norm" begin
            @test begin
                T = 50000
                z = 0
                for t in 1:T
                    bit = pop!(norm(v))
                    z += pos(bit) - neg(bit)
                end

                z / T
            end ≈ norm(v).value atol = 0.01
        end
    end
end


@testset "GenCB_03" begin
    env = CPLEX.Env()
    CPLEX.set_param!(env, "CPX_PARAM_THREADS", 1)
    model = CPLEX.Model(env, "LazyCB_test")
    #  max 0.5x + y
    # s.t. 0 <= x <= 2
    #      0 <= y <= 2
    CPLEX.add_var!(model, 0.5, 0, 2)
    CPLEX.add_var!(model, 1.0, 0, 2)
    CPLEX.set_sense!(model, :Max)
    CPLEX.set_vartype!(model, ['I', 'I'])
    function my_callback(cb_context::CPLEX.CallbackContext, context_id::Clong)
        if context_id == CPLEX.CPX_CALLBACKCONTEXT_CANDIDATE
            primal_sol = Vector{Float64}(undef, 2)
            objective = Ref{Float64}(0.0)
            CPLEX.cbgetcandidatepoint(
                cb_context,
                primal_sol,
                Cint(0),
                Cint(1),
                objective
            )
            @assert isapprox(
                0.5 * primal_sol[1] + primal_sol[2], objective[], atol=1e-6)
            if sum(primal_sol) > 3.0 + 1e-6
                CPLEX.cbrejectcandidate(
                    cb_context,
                    Cint(1),
                    Cint(2),
                    Cdouble[3.0],
                    Cchar['L'],
                    Cint[0],
                    Cint[0, 1],
                    Cdouble[1.0, 1.0]
                )
            end
        else
            error("Callback shold not be called from context_id $(context_id).")
        end
    end
    CPLEX.cbsetfunc(model, CPLEX.CPX_CALLBACKCONTEXT_CANDIDATE, my_callback)
    CPLEX.optimize!(model)
    sol = CPLEX.get_solution(model)
    @test sol ≈ [1, 2]
    @test CPLEX.get_objval(model) ≈ 2.5
    @test CPLEX.get_status(model) == :CPXMIP_OPTIMAL
end


@testset "NoiseGrid" begin

  using DiffEqNoiseProcess, DiffEqBase, Test

  t = 0:0.001:1
  grid = exp.(t)
  W = NoiseGrid(t,grid)

  dt = 0.1
  calculate_step!(W,dt,nothing,nothing)

  for i in 1:10
    accept_step!(W,dt,nothing,nothing)
  end

  W = NoiseGrid(t,grid)
  prob = NoiseProblem(W,(0.0,1.0))
  sol = solve(prob;dt=0.1)

  @test !sol.step_setup
  @test_throws ErrorException accept_step!(sol,dt,nothing,nothing)

  t = 0:0.001:1
  grid = [[exp.(ti) for i in 1:8] for ti in t]
  W = NoiseGrid(t,grid)
  prob = NoiseProblem(W,(0.0,1.0))
  sol = solve(prob;dt=0.1)

  dt = 0.001
  t = 0:dt:1
  brownian_values = cumsum([0;[sqrt(dt)*randn() for i in 1:length(t)-1]])
  W = NoiseGrid(t,brownian_values)

  dt = 0.001
  t = 0:dt:1
  brownian_values2 = cumsum([[zeros(8)];[sqrt(dt)*randn(8) for i in 1:length(t)-1]])
  W = NoiseGrid(t,brownian_values2)
  prob = NoiseProblem(W,(0.0,1.0))
  sol = solve(prob;dt=0.1)

  dt = 1//1000
  t = 0:dt:1
  W = NoiseGrid(t,brownian_values)
  prob_rational = NoiseProblem(W,(0,1))
  sol = solve(prob_rational; dt=1//10)
end


chebyshev_types = (Float32, Float64, LargeFloat)

function test_chebyshev_expansions(T)
    n = 10
    b = ChebyshevT{T}(n)
    if T == Float32 || T == Float64
        g = BasisFunctions.secondgrid(b)
        BasisFunctions.Test.test_generic_dict_transform(b, g)
        BasisFunctions.Test.test_generic_dict_transform(complex(b), g)
        @test convert(ChebyshevT{BigFloat}, ChebyshevT{T}(10)) isa ChebyshevT{BigFloat}
    end

    @test support(b) == ChebyshevInterval{T}()
end

function test_chebyshev_orthogonality()
    B = ChebyshevT(10)
    test_orthogonality_orthonormality(B, true, false, ChebyshevTWeight{Float64}())
    test_orthogonality_orthonormality(B, true, false, gauss_rule(B))
    test_orthogonality_orthonormality(B, true, false, discretemeasure(interpolation_grid(B)))
    test_orthogonality_orthonormality(B, true, false, discretemeasure(interpolation_grid(resize(B,length(B)+1))))
end

for T in chebyshev_types
    @testset "Chebyshev expansions ($T)" begin
        test_chebyshev_expansions(T)
    end
end

@testset "ChebyshevT orthogonality" begin
    test_chebyshev_orthogonality()
end


module TestTwoInertiasAndIdealGearWithMonteCarlo

using TinyModia
using DifferentialEquations
using ModiaPlot
using MonteCarloMeasurements
using Distributions
using Unitful

# The number of particles must be the same as for FloatType
const nparticles = 100
uniform(vmin,vmax) = StaticParticles(nparticles,Distributions.Uniform(vmin,vmax))

TwoInertiasAndIdealGearWithMonteCarlo = Model(
    J1 = 0.0025,
    J2 = uniform(50, 170),
    r  = 105,
    tau_max = 1,
    phi2 = Var(init = uniform(0.5,0.55)),
    w2 = Var(init = 0.0),
    equations = :[
        tau = if time < 1u"s"; tau_max elseif time < 2u"s"; 0 elseif time < 3u"s"; -tau_max else 0 end,

        # inertia1
        w1 = der(phi1),
        J1*der(w1) = tau - tau1,

        # gear
        phi1   = r*phi2,
        r*tau1 = tau2,

        # inertia2
        w2 = der(phi2),
        J2*der(w2) = tau2
    ]
)

twoInertiasAndIdealGearWithMonteCarlo = @instantiateModel(TwoInertiasAndIdealGearWithMonteCarlo,
                                                          FloatType = StaticParticles{Float64,nparticles})

simulate!(twoInertiasAndIdealGearWithMonteCarlo, Tsit5(), stopTime = 4.0,
          logParameters=true, logStates=true)

plot(twoInertiasAndIdealGearWithMonteCarlo, ["phi2", "w2"],
     MonteCarloAsArea=false)

end


module MarkovChains

# package code goes here
include("graphs.jl")
include("markov_chain.jl")
include("inftime_solve.jl")
include("fintime_solve.jl")
include("plot.jl")

struct CtmcSolution
    prob::Vector{Float64}
    cumtime::Vector{Float64}
end

export solve, state_prob, state_cumtime

function state_prob(sol::CtmcSolution, state::Integer)
    return sol.prob[state]
end

function state_cumtime(sol::CtmcSolution, state::Integer)
    return sol.cumtime[state]
end

function solve(chain, init_prob, time)
    n = state_count(chain)
    prob = Vector{Float64}(undef, n)
    cumtime = Vector{Float64}(undef, n)
    if isinf(time)
        sol = inftime_solve(chain, init_prob)
        for i in 1:n
            prob[i] = state_prob(sol, i)
            cumtime[i] = state_cumtime(sol, i)
        end
    else
        sol = fintime_solve(chain, init_prob, time)
        for i in 1:n
            prob[i] = state_prob(sol, i)
            cumtime[i] = state_cumtime(sol, i)
        end
    end
    CtmcSolution(prob, cumtime)
end

export Graphs

end # module


export OrthoNormal

doc"""
    OrthoNormal
"""
struct OrthoNormal
end

function (o::OrthoNormal)(::Type{T}, dim1::Int, dim2::Int) where T
    I = eye(dim2)
    lr = 0.1
    eps = 0.05 / (dim1+dim2)
    tries = 0
    while tries < 10
        Q = randn(dim1, dim2) / sqrt(dim2)
        for i = 1:100
            QTQmI = Q' * Q - I
            loss = sum(QTQmI .^ 2 / 2)
            Q2 = Q .^ 2
            a = abs.(Q2 .+ sum(Q2,1) .+ sum(Q2,2) - 1.0) + eps
            Q -= lr * Q * QTQmI ./ a
            if maximum(Q) > 1e6 || loss > 1e6 || isinf(loss)
                tries += 1
                lr /= 2.0
                break
            end
        end
        return Matrix{T}(Q)
    end
    #Q = xp.random.randn(input_size, output_size) / xp.sqrt(output_size)
    throw("Generation failed.")
end


#@testset "BigMat.jl" begin
    # tests for MPT_SVD, MPT_SVD, contract
    # vector to MPS
    v1      = [1.0; 4.5; 2; 7; 3; 7; 3.3; 10; 6; 8; 1.1; 4.6; 1; 6; 2; 7; 12; 4.5; 1; 0; 0; 3; 6; 2.1];
    v2      = [1.0; 4.5; 2; 7; 3; 7; 3.3; 10; 6; 8; 1.1; 4.6; 1; 6; 2; 7; 12; 4.5; 1; 0; 0; 3; 6; 2.1];
    M1      = randn(64,64)
    M2      = Symmetric(M1*M1')

    tt1,err = MPT_SVD(v1,[2 3 4],0.07);
    tt2     = TT_ALS(v2,[2 3 4],[1,2,4,1]);
    ttm1    = TTm_ALS(M1,[4 4 4; 4 4 4],[1,16,16,1]);
    ttm2    = TTm_ALS(M2,[4 4 4; 4 4 4],[1,16,16,1])
    ttm2    = TTm_ALS(Symmetric(M1*M1'),[4 4 4; 4 4 4],[1,3,3,1])

    @test (norm(tt1) - norm(tt2)) < 1e-13
    @test norm(mps2vec(tt1) - v1)/norm(v1) < 1e-13
    @test norm(mpo2mat(ttm1) - M1)/norm(M1) < 1e-13
    @test norm(mpo2mat(ttm2) - mpo2mat(ttm2)') < 1e-10

    # matrix to MPO
    M1 = v1*v1';
    M2 = v2*v2';
    mpo1,err = MPT_SVD(M1,[2 3 4; 2 3 4],0.0);
    mpo2 = MPT_ALS(M2,[2 3 4; 2 3 4],[1,4,16,1]);
    @test (norm(mpo1) - norm(mpo2)) < 1e-14

    @run mpo2mat(mpo1)
    @test norm(mpo2mat(mpo1) - M1)/norm(M1) < 1e-14
    @test norm(mpo1)-norm(mpo2mat(mpo1)) < 1e-14

    tt3 = MPT_SVD(v1,[2 3 4],0.07);
    @test norm(mps2vec(tt3)-v1)/norm(v1) <= 0.07

    # TT-ALS without orthogonalization needs to be checked. Wihtout pinv, the error is very high
    # maxiter + residulas need to be implemented

    # test of rktp2tn
    Ũ = [rand(50,40),rand(50,40),rand(50,40)];
    ϵ = 1.0;
    Umpo,err2 = rkrp2tn(Ũ,ϵ);
    Uapprox = mpo2mat(Umpo);
    U12 = zeros(2500,40);
    @inbounds @simd for j = 1:40
        @views kron!(U12[:,j],Ũ[1][:,j],Ũ[2][:,j])
    end
    Utrue = zeros(2500*50,40);
    @inbounds @simd for j = 1:40
        @views kron!(Utrue[:,j],U12[:,j],Ũ[3][:,j])
    end
    @test norm(Utrue-Uapprox)/norm(Utrue) < ϵ

    # get matrix entry from mpo2mat
    Ktest = zeros(24,24);
    linear = LinearIndices((1:2, 1:3, 1:4, 1:2, 1:3, 1:4));
    for i1 = 1:2
        for i2 = 1:3
            for i3 = 1:4
                for j1 = 1:2
                    for j2 = 1:3
                        for j3 = 1:4
                            tmp = mpo2[1][:,i1,j1,:]*mpo2[2][:,i2,j2,:]*mpo2[3][:,i3,j3,:];
                            Ktest[linear[i1,i2,i3,j1,j2,j3]] = tmp[1];
                        end
                    end
                end
            end
        end
    end
    @test norm(Ktest-M2)/norm(M2) < 1e-10

    # TT-SVD for symmetrical matrix
    X, y, f, K  = gensynthdata(1024,1,0.01,1.0,1.0)
    Kttm,Lttm,Uttm,S,err = MPT_SVD(Symmetric(K),[4 4 4 4 4; 4 4 4 4 4],0.1)
    k = mps2vec(Kttm)
    Kt = reshape(k,(4,4,4,4,4,4,4,4,4,4))
    Kt = permutedims(Kt,[1 2 3 4 5 10 9 8 7 6])
    Kapprox = reshape(Kt,(1024, 1024))
    norm(K[:]-Kt[:])/norm(K[:])
    rank(Kttm)

    L  = (Kttm[1]*Kttm[2]*Kttm[3]*Kttm[4]*Kttm[5])[1,:,:]
    L2 = mpo2mat(Lttm)
    norm(L*L' - K)/norm(K)
    norm(L2*L2' - K)/norm(K)
    U = mpo2mat(Uttm)
    norm(U*Diagonal(S)*U'-K)/norm(K)

    # test Hierarchical Tucker function
    tensor = rand(5,5,5);
    ranks  = [2,2,2];
    @run leaves2roottrunc(tensor,ranks)

#end


df = CSV.File("config/min_max_adj.csv") |> DataFrame!
const MINS = Array{Float64}(df[!, :min])
const MAXS = Array{Float64}(df[!, :adj_max])
const DIFF = MAXS .- MINS
const ON_INP = .~(Bool.(df[!, :off]))

function get_inputs(features::Array{Float64})
    inputs = (features .- MINS) ./ DIFF
    inputs = min.(max.(inputs, -1.0), 1.0)
    inputs[ON_INP]
end