"""
    ZeroOneNormalization(low::T, high::T)

This is a functor that normalizes a vector to be in the range [0,1]. Initial upper and lower bounds for each element are needed.

See also: [`PosNegNormalization`](@ref), [`GaussianNormalization`](@ref)

# Examples
```jldoctest
julia> low = [0.0, -1.0];

julia> high = [3.0, 0.5];

julia> nrm = ZeroOneNormalization(low, high);

julia> x = [1.0, 0.0];

julia> feats = nrm(x)
2-element Array{Float64,1}:
 0.3333333333333333
 0.6666666666666666

julia> y = zero(x);  # create buffer to prevent allocations

julia> feats = nrm(y, x);  # no allocation return

julia> feats = nrm(y, x, fit=true);  # update upper and lower bounds if x is outside the range

```
"""
struct ZeroOneNormalization{T} <: Any where {T}
    low::T
    high::T

    function ZeroOneNormalization(ranges)
        low = ranges[:, 1]
        high = ranges[:, 2]
        return new{typeof(low)}(low, high)
    end

    function ZeroOneNormalization(low::T, high::T) where {T}
        return new{T}(deepcopy(low), deepcopy(high))
    end
end

"""
    PosNegNormalization(low::T, high::T)

This is a functor that normalizes a vector to be in the range [-1,1]. Initial upper and lower bounds for each element are needed.

See also: [`ZeroOneNormalization`](@ref), [`GaussianNormalization`](@ref)

# Examples
```jldoctest
julia> low = [0.0, -1.0];

julia> high = [3.0, 0.5];

julia> nrm = PosNegNormalization(low, high);

julia> x = [1.0, 0.0];

julia> feats = nrm(x)
2-element Array{Float64,1}:
 -0.33333333333333337
  0.33333333333333326

julia> y = zero(x);  # create buffer to prevent allocations

julia> feats = nrm(y, x);  # no allocation return

julia> feats = nrm(y, x, fit=true);  # update upper and lower bounds if x is outside the range

```
"""
struct PosNegNormalization{T} <: Any where {T}
    low::T
    high::T

    function PosNegNormalization(ranges)
        low = ranges[:, 1]
        high = ranges[:, 2]
        return new{typeof(low)}(low, high)
    end

    function PosNegNormalization(low::T, high::T) where {T}
        return new{T}(deepcopy(low), deepcopy(high))
    end
end

"""
    GaussianNormalization(num_features::Int[, weight])

This is a functor that normalizes each element of a vector to be centered around a mean and have a variance of one.
OnlineStats.KahanVariance is used for computing and tracking the mean and variance. Likewise any OnlineStats.Weight can
be used to weigh each sample.

See also: [`ZeroOneNormalization`](@ref), [`PosNegNormalization`](@ref)

# Examples
```jldoctest
julia> nrm = GaussianNormalization(2, 1e-4);

julia> x = [1.0, 0.0];

julia> feats = nrm(x)
2-element Array{Float64,1}:
 1.0
 0.0

julia> y = zero(x);  # create buffer to prevent allocations

julia> feats = nrm(y, x);  # no allocation return

julia> feats = nrm(y, x, fit=true)  # update mean and variance using x
2-element Array{Float64,1}:
 0.9999
 0.0

```
"""
struct GaussianNormalization{T} <: Any where {T}
    v::T

    function GaussianNormalization(num_features::Int)
        v = num_features*KahanVariance()
        return new{typeof(v)}(v)
    end

    function GaussianNormalization(num_features::Int, weight::T) where {T<:Real}
        v = num_features*KahanVariance(weight=ExponentialWeight(weight))
        return new{typeof(v)}(v)
    end

    function GaussianNormalization(num_features::Int, weight::OnlineStats.Weight)
        v = num_features*KahanVariance(weight=weight)
        return new{typeof(v)}(v)
    end
end

function (f::ZeroOneNormalization)(x; fit=false)
    if fit
        @. f.low = min(x, f.low)
        @. f.high = max(x, f.high)
    end
    y = zero(x)
    @. y = (x - f.low) / (f.high - f.low)
    return y
end

function (f::ZeroOneNormalization)(buff::T, x::T; fit=false) where {T}
    if fit
        @. f.low = min(x, f.low)
        @. f.high = max(x, f.high)
    end
    @. buff = (x - f.low) / (f.high - f.low)
    return buff
end

function (f::PosNegNormalization)(x; fit=false)
    if fit
        @. f.low = min(x, f.low)
        @. f.high = max(x, f.high)
    end
    y = zero(x)
    T2 = eltype(x)
    a = convert(T2, 2.0)
    b = convert(T2, 1.0)
    @. y = a * (x - f.low) / (f.high - f.low) - b
    return y
end

function (f::PosNegNormalization)(buff::T, x::T; fit=false) where {T}
    if fit
        @. f.low = min(x, f.low)
        @. f.high = max(x, f.high)
    end
    T2 = eltype(x)
    a = convert(T2, 2.0)
    b = convert(T2, 1.0)
    @. buff = a * (x - f.low) / (f.high - f.low) - b
    return buff
end

function (f::GaussianNormalization)(x; fit=false)
    if fit
        fit!(f.v, x)
    end
    y = zero(x)
    T2 = eltype(x)
    @. y = (x - convert(T2, mean(f.v.stats))) / convert(T2, std(f.v.stats))
    return y
end

function (f::GaussianNormalization)(buff::T, x::T; fit=false) where {T}
    if fit
        fit!(f.v, x)
    end
    T2 = eltype(x)
    @. buff = (x - convert(T2, mean(f.v.stats))) / convert(T2, std(f.v.stats))
    return buff
end


# The following variables need to be defined before `include`-ing this file:
# `testset_name_suffix`
# `n_values`
# `k_values`
# `m_values`
function test_complex(::Type{TE}, m_values, k_values, n_values, testset_name_suffix) where {TE}
  T = Complex{TE}
  @time @testset "Matrix Multiply $T $(testset_name_suffix)" begin
    for n ∈ n_values
      for k ∈ k_values
        for m ∈ m_values
          A = rand(T, m, k);
          B = rand(T, k, n);
          b = rand(T, k);

          Are = real.(A);
          Bre = real.(B);
          bre = real.(b);

          A′ = permutedims(A)'
          B′ = permutedims(B)'
          AB = A * B;
          A′B = A′*B;
          AB′ = A*B′;
          A′B′= A′*B′;
          Ab = A*b;
          A′b = A′*b;

          AreB = Are*B;
          ABre = A*Bre;
          Areb = Are*b;
          Abre = A*bre;

          @info "" T n k m
          @test @time(Octavian.matmul(A, B)) ≈ AB
          @test @time(Octavian.matmul(A, Bre)) ≈ ABre
          @test @time(Octavian.matmul(Are, B)) ≈ AreB
          @test @time(Octavian.matmul(A′, B)) ≈ A′B
          @test @time(Octavian.matmul(A, B′)) ≈ AB′
          @test @time(Octavian.matmul(A′, B′)) ≈ A′B′

          @test @time(Octavian.matmul(A, b)) ≈ Ab
          @test transpose(@time(Octavian.matmul(transpose(b), transpose(A)))) ≈ Ab
          @test @time(Octavian.matmul(A, bre)) ≈ Abre
          @test @time(Octavian.matmul(Are, b)) ≈ Areb
          @test @time(Octavian.matmul(A′, b)) ≈ A′b
          @test transpose(@time(Octavian.matmul(transpose(b), transpose(A′)))) ≈ A′b

          @test @time(Octavian.matmul_serial(A, B)) ≈ AB
          @test @time(Octavian.matmul_serial(A, Bre)) ≈ ABre
          @test @time(Octavian.matmul_serial(Are, B)) ≈ AreB
          @test @time(Octavian.matmul_serial(A′, B)) ≈ A′B
          @test @time(Octavian.matmul_serial(A, B′)) ≈ AB′
          @test @time(Octavian.matmul_serial(A′, B′)) ≈ A′B′

          @test @time(Octavian.matmul_serial(A, b)) ≈ Ab
          @test transpose(@time(Octavian.matmul_serial(transpose(b), transpose(A)))) ≈ Ab
          @test @time(Octavian.matmul_serial(A, bre)) ≈ Abre
          @test @time(Octavian.matmul_serial(Are, b)) ≈ Areb
          @test @time(Octavian.matmul_serial(A′, b)) ≈ A′b
          @test transpose(@time(Octavian.matmul_serial(transpose(b), transpose(A′)))) ≈ A′b

          C = Matrix{T}(undef, n, m)'
          @test @time(Octavian.matmul!(C, A, B)) ≈ AB

          C1 = rand(T, m, n)
          C2 = copy(C1)
          α, β = T(1 - 2im), T(3 + 4im)
          @test @time(Octavian.matmul!(C1, A, B, α, β)) ≈ Octavian.matmul!(C2, A, B, α, β)
        end
      end
    end
  end
end

function matmul_pack_ab!(C, A, B)
  M, N = size(C); K = size(B,1)
  zc, za, zb = Octavian.zstridedpointer.((C,A,B))
  nspawn = min(Threads.nthreads(), Octavian.num_cores())
  GC.@preserve C A B begin
    if nspawn > 1
      threads, torelease = Octavian.PolyesterWeave.__request_threads((nspawn-1)%UInt32, Octavian.PolyesterWeave.worker_pointer(), nothing)
      @assert threads.i < Threads.nthreads()
      Octavian.matmul_pack_A_and_B!(
        zc, za, zb, Octavian.StaticInt{1}(), Octavian.StaticInt{0}(), M, K, N, threads,
        Octavian.W₁Default(), Octavian.W₂Default(), Octavian.R₁Default(), Octavian.R₂Default()
      )
      Octavian.PolyesterWeave.free_threads!(torelease)
    else
      Octavian.matmul_st_pack_A_and_B!(
        zc, za, zb, Octavian.StaticInt{1}(), Octavian.StaticInt{0}(), M, K, N,
        Octavian.W₁Default(), Octavian.W₂Default(), Octavian.R₁Default(), Octavian.R₂Default(), 0
      )
    end
  end
  C
end

function test_real(::Type{T}, m_values, k_values, n_values, testset_name_suffix) where {T}
  @time @testset "Matrix Multiply $T $(testset_name_suffix)" begin
    for n ∈ n_values
      for k ∈ k_values
        for m ∈ m_values
          A = rand(T, m, k)
          B = rand(T, k, n)
          b = rand(T, k)
          A′ = permutedims(A)'
          B′ = permutedims(B)'
          AB = A * B;
          Ab = A * b;
          @info "" T n k m
          @test @time(Octavian.matmul(A, B)) ≈ AB
          @test @time(Octavian.matmul(A′, B)) ≈ AB
          @test @time(Octavian.matmul(A, B′)) ≈ AB
          @test @time(Octavian.matmul(A′, B′)) ≈ AB
          @test @time(Octavian.matmul_serial(A, B)) ≈ AB
          @test @time(Octavian.matmul_serial(A′, B)) ≈ AB
          @test @time(Octavian.matmul_serial(A, B′)) ≈ AB
          @test @time(Octavian.matmul_serial(A′, B′)) ≈ AB
          @test @time(Octavian.matmul(A, b)) ≈ Ab
          @test @time(Octavian.matmul(A′, b)) ≈ Ab
          @test @time(Octavian.matmul(b', A'))' ≈ Ab
          @test @time(Octavian.matmul(b', A′'))' ≈ Ab
          @test @time(Octavian.matmul_serial(A, b)) ≈ Ab
          @test @time(Octavian.matmul_serial(A′, b)) ≈ Ab
          @test @time(Octavian.matmul_serial(b', A'))' ≈ Ab
          @test @time(Octavian.matmul_serial(b', A′'))' ≈ Ab
        end
      end
    end
    m = k = n = max(8Octavian.OCTAVIAN_NUM_TASKS[], 400)
    A = rand(T, m, k);
    B = rand(T, k, n);
    A′ = permutedims(A)';
    B′ = permutedims(B)';
    AB = A * B;
    @test matmul_pack_ab!(similar(AB), A, B) ≈ AB
    @test matmul_pack_ab!(similar(AB), A, B′) ≈ AB
    @test matmul_pack_ab!(similar(AB), A′, B) ≈ AB
    @test matmul_pack_ab!(similar(AB), A′, B′) ≈ AB
  end
end

@time @testset "zero-sized-matrices" begin
    @test Octavian.matmul_serial(randn(0,0), randn(0,0)) == zeros(0, 0)
    @test Octavian.matmul_serial(randn(2,3), randn(3,0)) == zeros(2, 0)
    @test Octavian.matmul_serial(randn(2,0), randn(0,2)) == zeros(2, 2)
    @test Octavian.matmul_serial!(ones(2,2),randn(2,0), randn(0,2), 1.0, 2.0) == ones(2, 2) .* 2
    @test Octavian.matmul(randn(0,0), randn(0,0)) == zeros(0, 0)
    @test Octavian.matmul(randn(2,3), randn(3,0)) == zeros(2, 0)
    @test Octavian.matmul(randn(2,0), randn(0,2)) == zeros(2, 2)
    @test Octavian.matmul!(ones(2,2),randn(2,0), randn(0,2), 1.0, 2.0) == ones(2, 2) .* 2
end


#=
# MNIST

We begin by importing the required packages.
We load MNIST via the MLDatasets.jl package.
=#

import Makie
import CairoMakie
import MLDatasets
import Flux
import RestrictedBoltzmannMachines as RBMs
import ConvolutionalRBMs as ConvRBMs
using Statistics: mean, var, std
using ValueHistories: MVHistory
using Random: bitrand
using RestrictedBoltzmannMachines: visible, hidden, weights, log_pseudolikelihood, transfer_sample
nothing #hide

#=
Useful function to plot MNIST digits.
=#

"""
    imggrid(A)

Given a four dimensional tensor `A` of size `(width, height, ncols, nrows)`
containing `width x height` images in a grid of `nrows x ncols`, this returns
a matrix of size `(width * ncols, height * nrows)`, that can be plotted in a heatmap
to display all images.
"""
function imggrid(A::AbstractArray{<:Any,4})
    (width, height, ncols, nrows) = size(A)
    return reshape(permutedims(A, (1,3,2,4)), width * ncols, height * nrows)
end

imggrid_border(A::AbstractArray{<:Any,4}, borderval = 1) = imggrid(bordered(A, borderval))

function bordered(A::AbstractArray{<:Any,4}, borderval = 1)
    return mapslices(A; dims=(1,2)) do img::AbstractMatrix
        [   fill(borderval, 1, size(img, 2) + 2);
            fill(borderval, size(img, 1)) img fill(borderval, size(img, 1));
            fill(borderval, 1, size(img, 2) + 2)
        ]
    end
end

#=
Load MNIST dataset.
=#

Float = Float32
train_x, train_y = MLDatasets.MNIST.traindata()
digit = 2
train_x = Array{Float}(train_x[:, :, train_y .== digit] .≥ 0.5)
train_y = train_y[train_y .== digit]
println(length(train_y), " training images (for digit = $digit)")
nothing #hide

# Reshape for convolutional input

train_x = reshape(train_x, 1, 28, 28, :) # channel dims, input dims, batch dims
nothing #hide

# Initialize the convolutional RBM.

rbm = ConvRBMs.BinaryConvRBM(Float, 1, 16, (15,15); pad=:same, pool=true)
RBMs.initialize!(rbm, train_x)
nothing #hide

# Pseudolikelihood before training

idx = rand(1:size(train_x)[end], 256)
mean(@time log_pseudolikelihood(rbm, train_x[:,:,:,idx]))

# Initialize training

batchsize = 256
optim = Flux.ADAM()
vm = transfer_sample(visible(rbm), falses(1, 28, 28, batchsize)) # fantasy chains
history = MVHistory()
nothing #hide

# Train!

@time for iter in 1:20
    ConvRBMs.pcd!(rbm, train_x; vm, history, batchsize, optim, epochs=5)
    lpl = log_pseudolikelihood(rbm, train_x[:, :, :, rand(1:size(train_x)[end], 1024)])
    push!(history, :lpl_ave, mean(lpl))
    push!(history, :lpl_std, std(lpl))
end
nothing #hide

# Plot of log-pseudolikelihood of trian data during learning.

fig = Makie.Figure(resolution=(600,300))
ax = Makie.Axis(fig[1,1], xlabel = "train time", ylabel="pseudolikelihood")
Makie.band!(ax, get(history, :lpl_ave)[1],
    get(history, :lpl_ave)[2] - get(history, :lpl_std)[2]/2,
    get(history, :lpl_ave)[2] + get(history, :lpl_std)[2]/2,
    color=:lightblue
)
Makie.lines!(ax, get(history, :lpl_ave)..., color=:blue)
fig

#=
Now let's generate some random RBM samples.
=#

nrows, ncols = 10, 15
nsteps = 1000
fantasy_F = zeros(nrows*ncols, nsteps)
fantasy_x = bitrand(1,28,28,nrows*ncols)
fantasy_F[:,1] .= RBMs.free_energy(rbm, fantasy_x)
@time for t in 2:nsteps
    fantasy_x .= RBMs.sample_v_from_v(rbm, fantasy_x)
    fantasy_F[:,t] .= RBMs.free_energy(rbm, fantasy_x)
end
nothing #hide

# Check equilibration of sampling

fig = Makie.Figure(resolution=(400,300))
ax = Makie.Axis(fig[1,1], xlabel="sampling time", ylabel="free energy")
fantasy_F_μ = vec(mean(fantasy_F; dims=1))
fantasy_F_σ = vec(std(fantasy_F; dims=1))
Makie.band!(ax, 1:nsteps, fantasy_F_μ - fantasy_F_σ/2, fantasy_F_μ + fantasy_F_σ/2)
Makie.lines!(ax, 1:nsteps, fantasy_F_μ)
fig

# Plot the resulting samples.

fig = Makie.Figure(resolution=(40ncols, 40nrows))
ax = Makie.Axis(fig[1,1], yreversed=true)
Makie.image!(ax, imggrid(reshape(fantasy_x, 28, 28, ncols, nrows)), colorrange=(1,0))
Makie.hidedecorations!(ax)
Makie.hidespines!(ax)
fig

# Plot the filters learned

wncols = 4; wnrows = 4
fig = Makie.Figure(resolution=(80wncols, 80wnrows))
ax = Makie.Axis(fig[1,1], yreversed=true)
Makie.image!(ax, imggrid_border(reshape(rbm.w ./ maximum(abs, rbm.w; dims=(2,3)), 15, 15, wncols, wnrows)))
Makie.hidedecorations!(ax)
Makie.hidespines!(ax)
fig


using BKMaxflow
using Test

# Julia API
weights, neighbors = create_graph(JuliaImpl{Float64,Int}, 4)

add_edge!(weights, neighbors, 1, 2, 1., 1.)
add_edge!(weights, neighbors, 1, 3, 2., 2.)
add_edge!(weights, neighbors, 2, 3, 3., 4.)
add_edge!(weights, neighbors, 2, 4, 5., 5.)
add_edge!(weights, neighbors, 3, 4, 6., 6.)

w = reshape(weights, 2, :)
flow, label = boykov_kolmogorov(1, 4, neighbors, w)
@test flow == 3
@test label[2] == label[3] == BK_T

# C API
# low-level
err = Ref{bk_error}(C_NULL)

g = bk_create_graph_double(2, 1, err[])

bk_add_node_double(g, 1, err[])
bk_add_node_double(g, 1, err[])

bk_add_tweights_double(g, 0, 1., 5., err[])
bk_add_tweights_double(g, 1, 2., 6., err[])

bk_add_edge_double(g, 0, 1, 3., 4., err[])

@test bk_maxflow_double(g, false, err[]) == 3

@test bk_what_segment_double(g, 0, err[]) == 1  # sink
@test bk_what_segment_double(g, 1, err[]) == 1  # sink

bk_delete_graph_double(g)

# high-level
g = create_graph(CImpl{Cdouble}, 2, 1)

add_node(CImpl{Cdouble}, g, 1)
add_node(CImpl{Cdouble}, g, 1)

add_tweights(CImpl, g, 0, 1., 5.)
add_tweights(CImpl, g, 1, 2., 6.)

add_edge(CImpl, g, 0, 1, 3., 4.)

@test maxflow(CImpl{Cdouble}, g) == 3

@test what_segment(CImpl{Cdouble}, g, 0) == 1  # sink
@test what_segment(CImpl{Cdouble}, g, 1) == 1  # sink

delete_graph(CImpl{Cdouble}, g)


### A Pluto.jl notebook ###
# v0.12.21

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ c7118f3c-209e-11eb-1f79-bdc2dd17aea0
begin
	using Revise
	using PlutoMustache
	using TropicalTensors
	using Random
	using LightGraphs
	using Compose, Viznet
	using SimpleTensorNetworks
	using HDF5
	Compose.set_default_graphic_size(14cm, 14cm)
end

# ╔═╡ 73cff25e-5434-11eb-0147-21ac9dbcb8b5
function rand_3regular_tn(::Type{T}, n; D=2) where T
	g = LightGraphs.random_regular_graph(n, 3)
	labels = 1:ne(g)
	arrays = [rand(T, fill(2, 3)...) for i=1:n]
	labels = [Int[] for i=1:n]
	for (k, e) in enumerate(edges(g))
		push!(labels[e.src], k)
		push!(labels[e.dst], k)
	end
	tensors = LabeledTensor.(arrays, labels)
	TensorNetwork(tensors)
end

# ╔═╡ 6b01ff60-217a-11eb-2a5c-13eb2716c727
md"number of nodes = $(@bind nnodes Slider(2:2:40, default=16))"

# ╔═╡ 2d5dd844-2136-11eb-176f-e53e2ac97d82
tnr = rand_3regular_tn(Float64, nnodes; D=2)

# ╔═╡ 8d7cb70c-5441-11eb-35cb-e551b9166fd6
md"contract the graph by a random order"

# ╔═╡ 1aabc318-2137-11eb-33b7-cfcec36e631f
@bind step Button(label="step")

# ╔═╡ 72d7474c-2137-11eb-030b-eb894e4b22dc
@regstate step

# ╔═╡ 09ea2b64-2137-11eb-187a-fdf0f5b882ca
@when step begin
	length(tnr.tensors) > 1 && SimpleTensorNetworks.contract_label!(tnr, rand(vcat([t.labels for t in tnr.tensors]...)))
	tnr
end

# ╔═╡ Cell order:
# ╠═c7118f3c-209e-11eb-1f79-bdc2dd17aea0
# ╠═73cff25e-5434-11eb-0147-21ac9dbcb8b5
# ╟─6b01ff60-217a-11eb-2a5c-13eb2716c727
# ╠═2d5dd844-2136-11eb-176f-e53e2ac97d82
# ╟─8d7cb70c-5441-11eb-35cb-e551b9166fd6
# ╟─1aabc318-2137-11eb-33b7-cfcec36e631f
# ╟─72d7474c-2137-11eb-030b-eb894e4b22dc
# ╠═09ea2b64-2137-11eb-187a-fdf0f5b882ca


export is_homogeneous,
    homogenize,
    uniquevar,
    Composition,
    expand,
    compose,
    validate,
    precondition,
    normalize_coefficients,
    linear_system

const MPPoly{T} = MP.AbstractPolynomialLike{T}
const MPPolys = AbstractVector{<:MP.AbstractPolynomialLike}
const WeightedVariable = Tuple{<:MP.AbstractVariable,Int}

##############
# COMPOSITION
##############

abstract type AbstractComposition end

"""
    Composition

A `Composition` is a composition of polynomial systems. This is the result of [`compose`](@ref).
"""
mutable struct Composition{T<:MP.AbstractPolynomialLike} <: AbstractComposition
    polys::Vector{Vector{T}} # polys = [f1, f2, f3] -> f1 ∘ f2 ∘ f3
end

Base.length(C::Composition) = length(C.polys[1])
Base.:(==)(C::Composition, D::Composition) = C.polys == D.polys

"""
    compose(g, f)::Composition

Compose the polynomial systems `g` and `f`.
You can also use the infix operator `∘` (written by \\circ).

```julia-repl
julia> @polyvar a b c x y z;
julia> g = [a * b * c];
julia> f = [x+y, y + z, x + z];
julia> expand(compose(g, f))
1-element Array{DynamicPolynomials.Polynomial{true,Int64},1}:
 x²y + x²z + xy² + 2xyz + xz² + y²z + yz²
```
"""
compose(g::Composition, f::Composition) = Composition([g.polys..., f.polys...])
compose(g::Composition, f::Composition, h::Composition...) =
    compose(Composition([g.polys..., f.polys...]), h...)
compose(g::Composition, fs::MPPolys...) = Composition([g.polys..., fs...])
compose(g::MPPolys, f::Composition) = Composition([g, f.polys...])
compose(g::MPPolys, f::Composition, h::Composition...) =
    compose(Composition([g, f.polys...]), h...)
compose(g::MPPolys, fs::MPPolys...) = Composition([g, fs...])

import Base: ∘
∘(g::Union{Composition,<:MPPolys}, f::MPPolys) = compose(g, f)
∘(g::Union{Composition,<:MPPolys}, f::MPPolys...) = compose(g, f...)
∘(g::Union{Composition,<:MPPolys}, f::Composition...) = compose(g, f...)

"""
    scale(C::Composition, λ)

Scale a composition by λ.
"""
scale(C::Composition, λ) = Composition([[C.polys[1] .* λ]; C.polys[2:end]])

function MP.subs(C::Composition, args::MP.AbstractSubstitution...)
    p = map(pᵢ -> MP.subs(pᵢ, args...), C.polys[end])
    Composition([C.polys[1:end-1]; [p]])
end

"""
    expand(C::Composition; parameters=nothing)

Expand the composition to the polynomial system is originally represents.

## Example
```julia-repl
julia> @polyvar a b c x y z;
julia> f = [a * b * c];
julia> g = [x+y, y + z, x + z];
julia> expand(f ∘ g)
1-element Array{DynamicPolynomials.Polynomial{true,Int64},1}:
 x²y + x²z + xy² + 2xyz + xz² + y²z + yz²
```
"""
function expand(C::Composition; parameters = nothing)
    g = nothing
    for f in reverse(C.polys)
        if g === nothing
            g = f
            continue
        end
        vars = variables(f, parameters)
        g = map(fᵢ -> MP.subs(fᵢ, vars => g), f)
    end
    g
end

"""
    validate(C::Composition; parameters=nothing)

Validates that the composition is well defined.
"""
function validate(C::Composition; parameters = nothing)
    g = nothing
    for f in reverse(C.polys)
        if g === nothing
            g = f
            continue
        end
        nvars = length(variables(f, parameters))
        if nvars !== length(g)
            return false
        end
        g = f
    end
    true
end

#############################
# HomogenizationInformation
#############################

"""
	HomogenizationInformation(;homvar=nothing, homvars=nothing, variable_groups=nothing)

This captures the information about the desired homogenization of the problem.
`homvar(s) = nothing` indicates that the problem should be homogenized.
`variable_groups` indicate the desired variable grouping for multi-homogeneous homotopies.

## Fields
* `homvars`
* `vargroups`
"""
struct HomogenizationInformation{N,T}
# one of the two fields is always not Nothing
    homvars::Union{Nothing,NTuple{N,T}}
    vargroups::Union{Nothing,NTuple{N,Vector{T}}}
end

function HomogenizationInformation(;
    homvar = nothing,
    homvars = nothing,
    variable_groups = nothing,
)
    if homvars === nothing && homvar !== nothing
        homvars = (homvar,)
        if variable_groups !== nothing && length(variable_groups) > 1
            error("You provided more than 1 variables group but only declared `homvar=$homvar`. Use instead `homvars=...`.")
        end
    end

    if variable_groups === nothing && homvars === nothing
        return nothing
    elseif variable_groups === nothing && length(homvars) > 1
        error("Currently we cannot find variable groups just from given `homvars`. Please provide explicit variable groups using `variable_groups=...`.")
    elseif variable_groups === nothing
        HomogenizationInformation(homvars, nothing)
    else
        HomogenizationInformation(homvars, tuple(collect.(variable_groups)...))
    end
end

function add_variable_groups(
    HI::HomogenizationInformation,
    F::MPPolys;
    parameters = nothing,
)
    if HI.vargroups === nothing
        if length(HI.homvars) == 1
            vargroups = (variables(F, parameters),)
            return HomogenizationInformation(HI.homvars, vargroups)
        else
            error("Cannot add variable groups")
        end
    end
    HI
end
add_variable_groups(::Nothing, F; parameters = nothing) = nothing

homvars(H::HomogenizationInformation) = H.homvars
homvars(::Nothing) = nothing

"""
    VariableGroups{N, Homvars, V}

A `VariableGroups` stores an `NTuple` of indices mapping to the indices of the original variables
of type `V`. `Homvars` is true if there have been homogenization variables declared.
"""
struct VariableGroups{N,Homvars,V<:Union{MP.AbstractVariable,Int}}
    variables::Vector{V}
    groups::NTuple{N,Vector{Int}}
    pull_back_mapping::Vector{Tuple{Int,Int}}
end

function VariableGroups(
    variables,
    groups::NTuple{N,Vector{Int}},
    dedicated_homvars::Bool,
) where {N}
    hom_vars = last.(groups)
    elements = vcat(groups...)
    pull_back_mapping = Tuple{Int,Int}[]
    proj_homvar_indices = cumsum([length(g) for g in groups])
    for i = 1:sum(length, groups)
        i ∉ hom_vars || continue
        l = 0
        for (gᵢ, group) in enumerate(groups), k in group
            l += 1
            if i == k
                push!(pull_back_mapping, (l, proj_homvar_indices[gᵢ]))
                break
            end
        end
    end
    VariableGroups{N,dedicated_homvars,eltype(variables)}(
        variables,
        groups,
        pull_back_mapping,
    )
end

"""
    VariableGroups(all_variables, variable_groups, homvars)

# Example
```julia
all_vars = @polyvar a b c x y z
variable_groups = ([c, b, a], [x, y, z])
homvars = (b, y)

vargroups = VariableGroups(all_vars, variable_groups, homvars)
vargroups.groups == ([3, 1, 2], [4, 6, 5])
vargroups.homvars == true
```
"""
function VariableGroups(
    all_variables,
    variable_groups::NTuple{N,<:Vector{<:Union{Int,MP.AbstractVariable}}},
    homvars::Union{Nothing,NTuple{N,<:Union{Int,MP.AbstractVariable}}} = nothing,
) where {N}
    groups = ntuple(N) do k
        variable_group = variable_groups[k]
        h = homvars === nothing ? 0 : homvars[k]
        group = Int[]
        homvar = 0
        for v in variable_group
            i = findfirst(w -> w == v, all_variables)
            if v == h
                homvar = i
            else
                push!(group, i)
            end
        end
        if homvar != 0
            push!(group, homvar)
        end
        group
    end
    VariableGroups(all_variables, groups, homvars !== nothing)
end
function VariableGroups(all_variables, homvar_info::HomogenizationInformation)
    if homvar_info.vargroups === nothing
        VariableGroups(all_variables, (collect(all_variables),), homvar_info.homvars)
    else
        VariableGroups(all_variables, homvar_info.vargroups, homvar_info.homvars)
    end
end

function VariableGroups(all_variables, homvar::MP.AbstractVariable)
    VariableGroups(all_variables, (collect(all_variables),), (homvar,))
end
function VariableGroups(all_variables, homvar::Nothing = nothing)
    VariableGroups(all_variables, (collect(all_variables),), homvar)
end
function VariableGroups(nvariables::Int, homvar::Nothing = nothing)
    VariableGroups(collect(1:nvariables), (collect(1:nvariables),), nothing)
end
function VariableGroups(nvariables::Int, homvar::Int)
    VariableGroups(collect(1:nvariables), (collect(1:nvariables),), (homvar,))
end
function VariableGroups(nvariables::Int, hominfo::HomogenizationInformation)
    if hominfo.vargroups === nothing
        VariableGroups(collect(1:nvariables), (collect(1:nvariables),), hominfo.homvars)
    else
        VariableGroups(collect(1:nvariables), hominfo.vargroups, hominfo.homvars)
    end
end

"""
	has_dedicated_homvars(variable_groups)

Returns `true` if there have been homogenization variables declared.
"""
has_dedicated_homvars(groups::VariableGroups{N,Homvars}) where {N,Homvars} = Homvars

"""
	projective_dims(variable_groups)

Returns the projective dimension of each variable group.
"""
function projective_dims(groups::VariableGroups)
    length.(groups.groups) .- 1
end

nvariables(G::VariableGroups) = sum(length.(G.groups))
Base.length(::VariableGroups{N}) where {N} = N

"""
	variable_groups(VG::VariableGroups)

Returns the variable groups.
"""
variable_groups(VG::VariableGroups) = map(g -> variables(VG)[g], VG.groups)

"""
	flattened_groups(VG::VariableGroups)

Returns the variable groups.
"""
flattened_variable_groups(VG::VariableGroups) = vcat(variable_groups(VG)...)

"""
	variables(VG::VariableGroups)

Returns all variables in their original order.
"""
variables(VG::VariableGroups) = VG.variables

"""
	ngroups(VG::VariableGroups)

Group the given variables in their corresponding groups.
"""
ngroups(VG::VariableGroups{M}) where {M} = M

function embed_projective(VG::VariableGroups{M}, v::AbstractVector) where {M}
    dims = projective_dims(VG)
    x = PVector(zeros(eltype(v), sum(dims) + M), dims)
    embed_projective!(x, VG, v)
end

function embed_projective!(
    x::PVector{<:Number,M},
    VG::VariableGroups{M},
    v::PVector{<:Number,M},
) where {M}
    x.data .= v.data
    x
end
function embed_projective!(
    x::PVector{<:Number,M},
    VG::VariableGroups{M},
    v::AbstractVector,
) where {M}
    if sum(projective_dims(VG)) == length(v)
        k = 1
        for group in VG.groups
            for i = 1:(length(group)-1)
                x[k] = v[group[i]]
                k += 1
            end
            x[k] = one(eltype(x))
            k += 1
        end
    else
        k = 1
        for group in VG.groups
            for i in group
                x[k] = v[i]
                k += 1
            end
        end
    end
    x
end

##############
# POLYNOMIALS
#############
function variables(
    polys::Union{MPPoly,MPPolys},
    params::AbstractVector{<:MP.AbstractVariable},
)
    setdiff!(variables(polys), params)
end
variables(C::Composition) = variables(C.polys[end])
variables(C::Composition, p::AbstractVector{<:MP.AbstractVariable}) =
    variables(C.polys[end], p)
variables(p::Union{MPPoly,MPPolys,Composition}, ::Nothing) = variables(p)

"""
    hasparameters(F, parameters)

Returns `true` if the parameters occur in F.
"""
function hasparameters(polys::Union{MPPoly,MPPolys}, parameters = nothing)
    isnothing(parameters) && return false
    vars = variables(polys)
    any(p -> p in vars, parameters)
end

"""
    nvariables(polys, parameters)

Returns the number of variables occuring in `polys`.
"""
nvariables(F::MPPolys, p::AbstractVector{<:MP.AbstractVariable}) = length(variables(F, p))
nvariables(C::Composition, p = nothing) = length(variables(C, p))

"""
    is_homogeneous(f::MP.AbstractPolynomialLike)

Checks whether `f` is homogeneous.

    is_homogeneous(f::MP.AbstractPolynomialLike, vars)

Checks whether `f` is homogeneous in the variables `vars` with possible weights.
"""
is_homogeneous(f::MP.AbstractPolynomialLike) = MP.mindegree(f) == MP.maxdegree(f)
function is_homogeneous(f::MP.AbstractPolynomialLike, variables)
    d_min, d_max = minmaxdegree(f, variables)
    d_min == d_max
end

"""
    is_homogeneous(F::Vector{MP.AbstractPolynomialLike}, variables)

Checks whether each polynomial in `F` is homogeneous in the variables `variables`.
"""
function is_homogeneous(F::MPPolys, variables)
    all(f -> is_homogeneous(f, variables), F)
end
function is_homogeneous(F::MPPolys, ::Nothing = nothing; parameters = nothing)
    if parameters !== nothing
        is_homogeneous(F, variables(F, parameters))
    else
        all(is_homogeneous, F)
    end
end
function is_homogeneous(
    F::MPPolys,
    hominfo::HomogenizationInformation;
    parameters = nothing,
)
    if hominfo.vargroups !== nothing
        all(vars -> is_homogeneous(F, vars), hominfo.vargroups)
    elseif parameters !== nothing
        is_homogeneous(F, variables(F, parameters))
    else
        all(is_homogeneous, F)
    end
end
function is_homogeneous(C::Composition, hominfo = nothing; kwargs...)
    homogeneous_degrees_helper(C; kwargs...) !== nothing
end
function is_homogeneous(C::Composition, hominfo::HomogenizationInformation; kwargs...)
    if hominfo.vargroups && length(hominfo.vargroups) > 1
        error("`variable_groups` and compositions are currently not supported.")
    end

    homogeneous_degrees_helper(C; kwargs...) !== nothing
end

function homogeneous_degrees_helper(C::Composition; parameters = nothing, weights = nothing)
    for f in reverse(C.polys)
        weights = homogeneous_degrees_helper(f, parameters = parameters, weights = weights)
        weights === nothing && return nothing
    end
    weights
end
function homogeneous_degrees_helper(F::MPPolys; parameters = nothing, weights = nothing)
    vars = variables(F, parameters)
    if !isnothing(weights)
        vars = zip(vars, weights)
    end
    degrees = Int[]
    for f in F
        mindeg, maxdeg = minmaxdegree(f, vars)
        if mindeg != maxdeg
            return nothing
        else
            push!(degrees, mindeg)
        end
    end
    degrees
end

"""
    degree(term::MP.AbstractTermLike, vars)

Compute the (weighted) degree of `f` in the variables `vars`.
"""
function degree(term::MP.AbstractTermLike, variables::Vector{<:MP.AbstractVariable})
    sum(MP.degree(term, v) for v in variables)
end
function degree(
    term::MP.AbstractTermLike,
    variables::NTuple{N,<:MP.AbstractVariable},
) where {N}
    sum(MP.degree(term, v) for v in variables)
end
function degree(term::MP.AbstractTermLike, weighted_variables)
    deg = 0
    for (v, w) in weighted_variables
        deg += w * MP.degree(term, v)
    end
    deg
end

"""
    minmaxdegree(f::MP.AbstractPolynomialLike, variables)

Compute the minimum and maximum (weighted total) degree of `f` with respect to the given variables.
"""
function minmaxdegree(f::MP.AbstractPolynomialLike, variables)
    d_min, d_max = typemax(Int), 0
    for term in MP.terms(f)
        d = degree(term, variables)
        d_min, d_max = min(d, d_min), max(d, d_max)
    end
    d_min, d_max
end

"""
    maxdegrees(F, parameters=nothing)

Computes the degrees of the polynomials of `F`.
"""
function maxdegrees(F::MPPolys; parameters = nothing, weights = nothing)
    vars = variables(F, parameters)
    if !isnothing(weights)
        vars = zip(vars, weights)
    end
    last.(minmaxdegree.(F, Ref(vars)))
end
function maxdegrees(C::Composition; parameters = nothing)
    weights = nothing
    for f in reverse(C.polys)
        weights = maxdegrees(f, parameters = parameters, weights = weights)
    end
    weights
end

"""
	multidegrees(F, variable_groups)

Computes the multi-degrees of the polynomials of `F` with respect to the given
variable groups. Returns an integer matrix, each column is the degree with respect
to the corresponding variable group and each row is one of the polynomials in the system.

## Example
```
julia> @polyvar x y;
julia> multidegrees([x*y-2, x^2-4], ([x], [y]))
2×2 Array{Int64,2}:
 1  2
 1  0
 ```
"""
multidegrees(F, vargroups) = [minmaxdegree(f, vars)[2] for vars in vargroups, f in F]
multidegrees(F, VG::VariableGroups) = multidegrees(F, variable_groups(VG))

"""
	support_coefficients(F::MPPolys, vars=variables(F))

Returns a tuple `(A, c)` where `A` is the vector of supports of `F` and `c` is the vector
of corresponding coefficients.
"""
function support_coefficients(F::MPPolys, vars = variables(F))
    terms₁ = MP.terms(F[1])
    c₁ = float.(MP.coefficient.(terms₁))
    A₁ = [Int32(MP.degree(term, var)) for var in vars, term in terms₁]
    A = [A₁]
    c = [c₁]
    for i = 2:length(F)
        termsᵢ = MP.terms(F[i])
        cᵢ = float.(MP.coefficient.(termsᵢ))
        Aᵢ = [Int32(MP.degree(term, var)) for var in vars, term in termsᵢ]
        push!(A, Aᵢ)
        push!(c, cᵢ)
    end
    A, c
end

"""
    compute_numerically_degrees(F::AbstractSystem)

This tries to compute numerically the degrees of `F`. If successfull
this returns a vector containing the degrees. If not `nothign` is returned.
"""
function compute_numerically_degrees(F)
    n, N = size(F)
    # The number of variables match, but it still cannot be homogeneous.
    # We evaluate the system with y:=rand(N) and 2y. If homogeneous then the output
    # scales accordingly to the degrees which we can obtain by taking logarithms.
    x = PVector(rand(ComplexF64, N))
    system_cache = cache(F, x)
    y = evaluate(F, x, system_cache)
    LinearAlgebra.rmul!(x, 2)
    y2 = evaluate(F, x, system_cache)

    degrees = Int[]
    for (y2ᵢ, yᵢ) in zip(y2, y)
        # y2ᵢ = 2^dᵢ yᵢ
        float_dᵢ = log2(abs(y2ᵢ / yᵢ))
        dᵢ = round(Int, float_dᵢ)
        if abs(dᵢ - float_dᵢ) > 1e-10
            return nothing
        end
        push!(degrees, dᵢ)
    end
    degrees
end

"""
    uniquevar(f::MP.AbstractPolynomialLike, tag=:x0)
    uniquevar(F::MPPolys, tag=:x0)

Creates a unique variable.
"""
uniquevar(f::MP.AbstractPolynomialLike, tag = :x0) = MP.similarvariable(f, gensym(tag))
uniquevar(F::MPPolys, tag = :x0) = uniquevar(F[1], tag)
uniquevar(C::Composition, tag = :x0) = uniquevar(C.polys[1][1], tag)

"""
    homogenize(f::MP.AbstractPolynomial, variable=uniquevar(f))

Homogenize the polynomial `f` by using the given variable `variable`.

    homogenize(F::Vector{<:MP.AbstractPolynomial}, variable=uniquevar(F))

Homogenize each polynomial in `F` by using the given variable `variable`.

    homogenize(f::MP.AbstractPolynomial, v::Vector{<:MP.AbstractVariable}, variable=uniquevar(f))

Homogenize the variables `v` in the polynomial `f` by using the given variable `variable`.

    homogenize(F::Vector{<:MP.AbstractPolynomial}, v::Vector{<:MP.AbstractVariable}, variable=uniquevar(F))

Homogenize the variables `v` in each polynomial in `F` by using the given variable `variable`.
"""
function homogenize(f::MP.AbstractPolynomialLike, var = uniquevar(f); parameters = nothing)
    homogenize(f, variables(f, parameters), var)
end
function homogenize(f::MP.AbstractPolynomialLike, vars::Vector, var = uniquevar(f))
    _, d_max = minmaxdegree(f, vars)
    MP.polynomial(map(f) do t
        d = degree(t, vars)
        var^(d_max - d) * t
    end)
end
function homogenize(f::MP.AbstractPolynomialLike, hominfo::HomogenizationInformation)
    g = f
    for (vargroup, homvar) in zip(hominfo.vargroups, hominfo.homvars)
        g = homogenize(g, vargroup, homvar)
    end
    g
end

function homogenize(F::MPPolys, var = uniquevar(F); parameters = nothing)
    homogenize(F, variables(F, parameters), var)
end
function homogenize(F::MPPolys, variables::Vector, var::MP.AbstractVariable = uniquevar(F))
    map(f -> homogenize(f, variables, var), F)
end
function homogenize(F::MPPolys, hominfo::HomogenizationInformation)
    homogenize.(F, Ref(hominfo))
end
function homogenize(
    C::Composition,
    var = uniquevar(C.polys[1]);
    parameters = nothing,
    weights = nothing,
)
    polys = map(length(C.polys):-1:1) do k
        f̄, weights = homogenize_degrees(
            C.polys[k],
            var;
            parameters = parameters, weights = weights,
        )
        if k > 1
            push!(f̄, var)
            push!(weights, 1)
        end
        f̄
    end
    Composition(reverse!(polys))
end

"""
    homogenize_degree(f::MPPoly, variables, var)

Homogenize `f` by using the variable `v`. Returns the homogenized polynomial
and its degree.
"""
function homogenize_degree(f::MPPoly, variables, var::MP.AbstractVariable)
    _, d = minmaxdegree(f, variables)
    MP.polynomial(map(t -> t * var^(d - degree(t, variables)), f)), d
end

"""
    homogenize_degrees(F::MPPolys, var; parameters=nothing, weights=nothing)
    homogenize_degrees(F::MPPolys, variables, var)

Homogenize the system `F` using the variable `var`. Returns the homogenized system
and the degrees of each polynomial.
"""
function homogenize_degrees(F::MPPolys, variables, var::MP.AbstractVariable)
    F̄ = similar(F)
    degrees = Int[]
    for (i, f) in enumerate(F)
        f̄, d = homogenize_degree(f, variables, var)
        F̄[i] = f̄
        push!(degrees, d)
    end

    F̄, degrees
end
function homogenize_degrees(
    F::MPPolys,
    var::MP.AbstractVariable;
    parameters = nothing,
    weights = nothing,
)
    allvars = variables(F, parameters)
    if weights !== nothing
        homogenize_degrees(F, zip(allvars, weights), var)
    else
        homogenize_degrees(F, allvars, var)
    end
end

"""
    remove_zeros!(F::Vector{<:MP.AbstractPolynomialLike})

Remove zero polynomials from the given system `F`.
"""
function remove_zeros!(F::MPPolys)
    filter!(!iszero, F)
    F
end
function remove_zeros!(C::Composition)
    filter!(!iszero, C.polys[1])
    C
end

"""
    has_constant_polynomial(F::Vector{<:MP.AbstractPolynomialLike})

Returns true if the system `F` contains a constant polynomial.
"""
function has_constant_polynomial(F::MPPolys)
    for f in F
        MP.nterms(f) == 1 && MP.isconstant(first(MP.terms(f))) && return true
    end
    false
end
has_constant_polynomial(C::Composition) = has_constant_polynomial(C.polys[1])

permute(F::MPPolys, perm) = F[perm]
permute(C::Composition, perm) = Composition([[C.polys[1][perm]]; C.polys[2:end]])

"""
    check_zero_dimensional(F::Vector{<:MP.AbstractPolynomial})

Check that the given polynomial system can have zero dimensional components.
"""
function check_zero_dimensional(F::Union{MPPolys,Composition})
    N = nvariables(F)
    n = length(F)

    if n ≥ N || (n == N - 1 && is_homogeneous(F))
        return nothing
    end
    error("The input system will not result in a finite number of solutions.")
end

"""
    homogenize_if_necessary(F::MPPolys)

Homogenizes the system `F` if necessary and returns the (new) system `F` its variables
and a subtype of [`AbstractHomogenization`] indicating whether it was homegenized.
If it was homogenized and no then the new variable is the **first one**.
"""
function homogenize_if_necessary(
    F::Union{MPPolys,Composition},
    hominfo::Union{Nothing,HomogenizationInformation};
    vars = nothing,
    parameters = nothing,
)
    if vars === nothing
        vars = variables(F, parameters)
    end

    # This fills in the simple variable group (allvars,)
    hominfo = add_variable_groups(hominfo, F; parameters = parameters)

    if is_homogeneous(F, hominfo; parameters = parameters)
        vargroups = VariableGroups(vars, hominfo)
        F, vargroups, homvars(hominfo)
    elseif hominfo === nothing
        homvar = uniquevar(F)
        push!(vars, homvar)

        F′ = homogenize(F, homvar; parameters = parameters)
        vargroups = VariableGroups(vars, homvar)

        F′, vargroups, (homvar,)
    else
        if hominfo.homvars !== nothing
            error("Input system is not homogeneous although `homvars=$(hominfo.homvars)` was passed.")
        end

        # We create a new variable for each variable group to homogenize the system
        new_homvars = map(_ -> uniquevar(F), hominfo.vargroups)
        push!(vars, new_homvars...)
        new_vargroups = map(hominfo.vargroups, new_homvars) do group, v
            vcat(group, v)
        end
        hominfo = HomogenizationInformation(new_homvars, new_vargroups)

        F′ = homogenize(F, hominfo)
        vargroups = VariableGroups(vars, hominfo)

        F′, vargroups, new_homvars
    end
end

"""
    classify_system(F, vargroups::VariableGroups; affine_tracking=false)

Returns a symbol indicating whether `F` is `:square`, `:overdetermined` or `:underdetermined`.
"""
function classify_system(F, vargroups::VariableGroups; affine_tracking = false)
    n = length(F) - nvariables(vargroups)
    if !affine_tracking
        n += ngroups(vargroups)
    end

    if n == 0
        :square
    elseif n > 0
        :overdetermined
    else
        :underdetermined
    end
end

exponent(term::MP.AbstractTermLike, vars) = [MP.degree(term, v) for v in vars]

"""
    weyldot(f::Polynomial, g::Polynomial)

Compute the [Bombieri-Weyl dot product](https://en.wikipedia.org/wiki/Bombieri_norm).
Note that this is only properly defined if `f` and `g` are homogeneous.

    weyldot(f::Vector{Polynomial}, g::Vector{Polynomial})

Compute the dot product for vectors of polynomials.
"""
function weyldot(
    f::MP.AbstractPolynomialLike{T},
    g::MP.AbstractPolynomialLike{S},
    vars = variables([f, g]),
) where {T,S}
    if f === g
        return sum(f) do term
            abs2(MP.coefficient(term)) / multinomial(exponent(term, vars))
        end
    end
    result = zero(promote_type(T, S))
    for term_f in f
        c_f = MP.coefficient(term_f)
        exp_f = exponent(term_f, vars)
        for term_g in g
            c_g = MP.coefficient(term_g)
            exp_g = exponent(term_g, vars)
            if exp_f == exp_g
                result += (c_f * conj(c_g)) / multinomial(exp_f)
                break
            end
        end
    end
    result
end
function weyldot(F::MPPolys, G::MPPolys, vars = variables(F))
    sum(fg -> weyldot(fg..., vars), zip(F, G))
end

"""
    weylnorm(F, vars=variables(F))

Compute the [Bombieri-Weyl norm](https://en.wikipedia.org/wiki/Bombieri_norm).
Note that this is only properly defined if `f` is homogeneous.
"""
weylnorm(F::MPPolys, vars = variables(F)) = √(sum(f -> weyldot(f, f, vars), F))
weylnorm(f::MPPoly, vars = variables(f)) = √(weyldot(f, f, vars))

"Computes the multinomial coefficient (|k| \\over k)"
function multinomial(k::Vector{Int})
    s = 0
    result = 1
    for i in k
        s += i
        result *= binomial(s, i)
    end
    result
end

"""
	precondition(f::Union{MPPolys,Composition}, variables=variables(f))

Rescale the equations and variables of `f` following the heuristic described in [1].
Returns the scaled system and a vector of scaling factors of the variables.

[1]: HOM4PS-2.0: a software package for solving polynomial systems by the polyhedral homotopy continuation method,
	Lee, T.L., Li, T.Y. & Tsai, C.H. Computing (2008) 83: 109.
	https://doi.org/10.1007/s00607-008-0015-6
"""
function precondition(f::Composition, variables = variables(f))
    vars, equations = preconditioning_factors(expand(f), variables)
    y = vars .* variables
    scale(subs(f, variables => y), equations), vars
end

function precondition(f::MPPolys, variables = variables(f))
    vars, equations = preconditioning_factors(f, variables)
    y = vars .* variables
    map(equations, f) do c, fᵢ
        c * MP.subs(fᵢ, variables => y)
    end, vars
end
function preconditioning_factors(f::MPPolys, vars)
    nvars = length(vars)
    m = sum(nᵢ -> nᵢ + div(nᵢ * (nᵢ - 1), 2), MP.nterms.(f))
    n = nvars + length(f)
    A, b = zeros(Int, m, n), zeros(Float64, m)

    row = 1
    for (i, fᵢ) in enumerate(f)
        termsᵢ = MP.terms(fᵢ)
        ntermsᵢ = length(termsᵢ)
        blockᵢ_start = row - 1
        for t in termsᵢ
            for (j, vⱼ) in enumerate(vars)
                A[row, j] = MP.degree(t, vⱼ)
            end
            A[row, nvars+i] = 1
            b[row] = -log10(convert(Float64, abs(MP.coefficient(t))))
            row += 1
        end

        for k = 1:ntermsᵢ, l = (k+1):ntermsᵢ
            for j = 1:nvars
                A[row, j] = A[blockᵢ_start+k, j] - A[blockᵢ_start+l, j]
            end
            b[row] = b[blockᵢ_start+k] - b[blockᵢ_start+l]
            row += 1
        end
    end
    c = ones(n)
    try
        c = LinearAlgebra.qr(A, Val(true)) \ b
    catch e
        if !isa(e, LinearAlgebra.LAPACKException)
            rethrow(e)
        end
    end
    (variables = exp10.(c[1:nvars]), equations = exp10.(c[nvars+1:end]))
end

"""
	normalize_coefficients(f::Union{MPPolys,Composition})

Rescale the equations of `f` such that each coefficient has absolute value between 0 and 1.
"""
function normalize_coefficients(f::Composition)
    Λ = maximum.(abs, MP.coefficients.(expand(f)))
    scale(f, Λ)
end
function normalize_coefficients(f::MPPolys)
    map(f) do fᵢ
        cmax = maximum(abs, MP.coefficients(fᵢ))
        map(MP.terms(fᵢ)) do t
            c = MP.coefficient(t)
            m = MP.monomial(t)
            c / cmax * m
        end |> MP.polynomial
    end
end

"""
    linear_system(f::Vector{<:MP.AbstractPolynomialLike})

Given a polynomial system which represents a linear system ``Ax=b`` return
`A` and `b`. If `f ` is not a linear system `nothing` is returned.

### Example
```
julia> @polyvar x y z;
julia> f = [2x+3y+z+5, -1x+2y+4z-2];
julia> A, b = linear_system(f);
julia> A
2×3 Array{Int64,2}:
  2  3  1
 -1  2  4

julia> b
2-element Array{Int64,1}:
 -5
  2
```
"""
function linear_system(f::Vector{<:MP.AbstractPolynomialLike}, vars = variables(f))
    n = length(vars)
    A = zeros(MP.coefficienttype(f[1]), length(f), n)
    b = zeros(eltype(A), length(f))
    for (i, fᵢ) in enumerate(f)
        for t in MP.terms(fᵢ)
            constant = true
            for (j, v) in enumerate(vars)
                d = MP.degree(t, v)
                d ≤ 1 || return nothing

                if d == 1
                    A[i, j] = MP.coefficient(t)
                    constant = false
                    break
                end
            end
            if constant
                b[i] = -MP.coefficient(t)
            end
        end
    end
    A, b
end

"""
    monomials(vars::Vector{<:PolyVar}, d; homogeneous::Bool = false)

Create all monomials of a given degree.

```
julia> @polyvar x y
(x, y)

julia> monomials([x,y], 2)
6-element Array{Monomial{true},1}:
 x²
 xy
 y²
 x
 y
 1

julia> monomials([x,y], 2; homogeneous = true)
3-element Array{Monomial{true},1}:
 x²
 xy
 y²
```
"""
function monomials(
    vars::AbstractVector{<:DynamicPolynomials.PolyVar},
    d::Int;
    homogeneous::Bool = false,
)
    n = length(vars)
    if homogeneous
        pred = x -> sum(x) == d
    else
        pred = x -> sum(x) ≤ d
    end
    exps = collect(Iterators.filter(pred, Iterators.product(Iterators.repeated(0:d, n)...)))
    sort!(exps, lt = td_order, rev = true)
    map(exps) do exp
        prod(i -> vars[i]^exp[i], 1:n)
    end
end
function td_order(x, y)
    sx = sum(x)
    sy = sum(y)
    sx == sy ? x < y : sx < sy
end

################
## ModelKit ##
################

polyvar(v::ModelKit.Variable) = DynamicPolynomials.PolyVar{true}(string(v.name))
function maxdegrees(F::ModelKit.System; parameters = F.parameters)
    variables = setdiff!([F.variables; F.parameters], parameters)
    vars = polyvar.(variables)
    params = polyvar.(parameters)
    P = evaluate(F.expressions, variables => vars, parameters => params)
    maxdegrees(P; parameters = params)
end


using QuickPOMDPs
using POMDPModelTools
using Distributions

#=
The simple 1-D light dark problem from https://arxiv.org/pdf/1709.06196v6.pdf, section 5.2, or https://slides.com/zacharysunberg/defense-4#/39
=#

r = 60
light_loc = 10

simple_lightdark = QuickPOMDP(
    states = -r:r+1,                  # r+1 is a terminal state
    actions = [-10, -1, 0, 1, 10],
    discount = 0.95,
    isterminal = s -> !(s in -r:r),
    obstype = Float64,

    transition = function (s, a)
        if a == 0
            return SparseCat(r+1, 1.0)
        else
            return SparseCat(clamp(s+a, -r, r), 1.0)
        end
    end,

    observation = (s, a, sp) -> Normal(sp, abs(sp - light_loc) + 0.0001),

    reward = function (s, a, sp, o)
        if a == 0
            return s == 0 ? 100 : -100
        else
            return -1.0
        end
    end,

    initialstate_distribution = POMDPModelTools.Uniform(div(-r,2):div(-r,2))
)


module Tst

using Plots
using JacobiDavidson
using LinearMaps
using LinearAlgebra
using SparseArrays

import LinearAlgebra: ldiv!

function myA!(y::AbstractVector{T}, x::AbstractVector{T}) where {T<:Number}
  for i = 1 : length(x)
    @inbounds y[i] = sqrt(i) * x[i]
  end
end

function myB!(y::AbstractVector{T}, x::AbstractVector{T}) where {T<:Number}
  for i = 1 : length(x)
    @inbounds y[i] = x[i] / sqrt(i)
  end
end

struct SuperPreconditioner{numT <: Number}
  target::numT
end

function ldiv!(p::SuperPreconditioner{T}, x::AbstractVector{T}) where {T<:Number}
  for i = 1 : length(x)
    @inbounds x[i] = x[i] * sqrt(i) / (i - p.target)
  end
  return x
end

function ldiv!(y::AbstractVector{T}, p::SuperPreconditioner{T}, x::AbstractVector{T}) where {T<:Number}
  for i = 1 : length(x)
    @inbounds y[i] = x[i] * sqrt(i) / (i - p.target)
  end
  return y
end

function another_example(; n = 1000, target = Near(31.0 + 0.1im))
  A = LinearMap{Float64}(myA!, n; ismutating = true)
  B = LinearMap{Float64}(myB!, n; ismutating = true)
  P = SuperPreconditioner(target.τ)

  schur2, residuals2 = jdqz(
    A, B,
    solver = GMRES(n, iterations = 3),
    preconditioner = P,
    testspace = Harmonic,
    target = target,
    pairs = 5,
    tolerance = 1e-9,
    subspace_dimensions = 10:20,
    max_iter = 100,
    verbosity = 1
  )

  schur, residuals = jdqz(
    A, B,
    solver = BiCGStabl(n, max_mv_products = 10, l = 2),
    preconditioner = P,
    testspace = Harmonic,
    target = target,
    pairs = 5,
    tolerance = 1e-9,
    subspace_dimensions = 10:20,
    max_iter = 100,
    verbosity = 1
  )

  plot(residuals, yscale = :log10, label = "BiCGStab", marker = :x)
  plot!(residuals2, yscale = :log10, label = "GMRES", marker = :x)
end

function generalized(; n = 1_000, target = Near(0.5 + 0.1im))
  A = 2I + sprand(ComplexF64, n, n, 1 / n)
  B = 2I + sprand(ComplexF64, n, n, 1 / n)

  values = eigvals(Matrix(A), Matrix(B))

  @time schur, residuals = jdqz(
    A, B,
    solver = GMRES(n, iterations = 10),
    target = target,
    pairs = 10,
    subspace_dimensions = 10:15,
    max_iter = 300,
    verbosity = 1
  )

  found = schur.alphas ./ schur.betas

  p1 = scatter(real(values), imag(values), label = "eig")
  scatter!(real(found), imag(found), marker = :+, label = "jdqz")
  scatter!([real(target.τ)], [imag(target.τ)], marker = :star, label = "Target")

  p2 = plot(residuals, yscale = :log10)

  p1, p2
end

function test_harmonic_2(n = 1000; τ = Near(1.52 + 0.0im))
  A = spdiagm(
    -1 => fill(-1.0, n - 1),
     0 => fill(2.0, n),
     1 => fill(-1.0, n - 1)
  )

  B = spdiagm(0 => fill(1.0, n))

  values = eigvals(Matrix(A), Matrix(B))

  schur, residuals = jdqz(
    A,
    B,
    solver = Exact(),
    pairs = 10,
    subspace_dimensions = 10:20,
    max_iter = 300,
    tolerance = 1e-5,
    target = τ
  )

  found = schur.alphas ./ schur.betas

  p1 = scatter(real(values), imag(values), ylims = (-1, 1))
  scatter!(real(found), imag(found), marker = :+)

  p2 = plot(residuals, yscale = :log10)

  p1, p2
end

function test_harmonic(; n = 500, τ = Near(0.01 + 0.02im))
  A = spdiagm(
    -1 => fill(-1.0, n - 1),
     0 => fill(2.0, n),
     1 => fill(-1.0, n - 1)
  )

  schur, ritz_hist, conv_hist, residuals = jdqr(
    A,
    solver = BiCGStabl(size(A, 1), max_mv_products = 30, l = 2),
    pairs = 10,
    subspace_dimensions = 5:10,
    max_iter = 300,
    tolerance = 1e-5,
    target = τ,
    verbosity = 1
  )

  λs = real(eigvals(Matrix(A)))

  # Total number of iterations
  iterations = length(ritz_hist)

  @show iterations
  @show schur.values

  # Plot the target as a horizontal line
  p = plot([0.0, iterations + 1.0], [real(τ.τ), real(τ.τ)], linewidth=5, legend = :none, layout = (2, 1))

  # Plot the actual eigenvalues
  scatter!(fill(iterations + 1, n), λs, xlabel = "Iteration", ylims = (minimum(λs) - 2.0, maximum(λs) + 2.0), marker = (:diamond, :black), subplot = 1)

  # Plot the approximate eigenvalues per iteration
  for (k, (ritzvalues, eigenvalues)) = enumerate(zip(ritz_hist, conv_hist))
    scatter!(k * ones(length(ritzvalues)), real(ritzvalues), marker = (:+, :green), subplot = 1)
    scatter!(k * ones(length(eigenvalues)), real(eigenvalues), marker = (:hexagon, :yellow), subplot = 1)
  end

  plot!(residuals, xlabel = "Iteration", label = "Residual", xticks = 0.0 : 50.0 : 350.0, yaxis = :log10, subplot = 2)

  p
end
end


using SpecialFunctions
"""
    gammaf()
Variables: 0.5 <= x <= 5;
Objective: Minimize gamma(x)
"""

function gammaf()
    best_known_objective = gamma(1.4616321)
    cplex_optimizer =
        JuMP.optimizer_with_attributes(CPLEX.Optimizer, "CPX_PARAM_SCRIND" => 0)
    error_tolerances = [NaN64, 1e-3]
    # base partition holds the inflection points and the end-points
    milp_statistics = zeros(length(error_tolerances), 5)
    lp_statistics = zeros(length(error_tolerances), 5)

    base_partition = [0.5, 5]
    # functions
    f = x -> gamma(x)

        # solve MILP for various error tolerances
        for i = 1:length(error_tolerances)
            error_tolerance = error_tolerances[i]
            # construct MILP relaxations
            milp = Model(cplex_optimizer)
            var, x_index, y_index = add_milp_relaxation(
                    f,
                    base_partition,
                    error_tolerance,
                    milp,
                    name = "1",
                )
            @objective(milp, Min, var[y_index])
            _, solve_time, solve_bytes_alloc, sec_in_gc = @timed optimize!(milp)
            relaxation_objective = objective_value(milp)
            @assert relaxation_objective ≤ best_known_objective
            relative_gap =
                abs(best_known_objective - relaxation_objective) / (abs(relaxation_objective) + 1e-4)
            milp_statistics[i, :] = [
                error_tolerance,
                best_known_objective,
                relaxation_objective,
                relative_gap,
                round(solve_time, digits = 2),
            ]
        end

        for i = 1:length(error_tolerances)
            error_tolerance = error_tolerances[i]
            # construct LP relaxations
            lp = Model(cplex_optimizer)
            var, x_index, y_index = add_lp_relaxation(
                    f,
                    base_partition,
                    error_tolerance,
                    lp,
                    name = "1",
                )
            @objective(lp, Min, var[y_index])
            _, solve_time, solve_bytes_alloc, sec_in_gc = @timed optimize!(lp)
            relaxation_objective = objective_value(lp)
            @assert relaxation_objective ≤ best_known_objective
            relative_gap =
                abs(best_known_objective - relaxation_objective) / abs(relaxation_objective)
            lp_statistics[i, :] = [
                error_tolerance,
                best_known_objective,
                relaxation_objective,
                relative_gap,
                round(solve_time, digits = 2),
            ]
        end

        open("csv/gammaf_milp.csv", "w") do io
            writedlm(io, ["error_tol" "g_opt" "relaxation" "relative_gap" "time"], ',')
            writedlm(io, milp_statistics, ",")
        end

        open("csv/gammaf_lp.csv", "w") do io
            writedlm(io, ["error_tol" "g_opt" "relaxation" "relative_gap" "time"], ',')
            writedlm(io, lp_statistics, ",")
        end

end


struct InstrumentalResponse{T<:Real}
    amp::Vector{T}
    width::Vector{T}    #in pixels
    location::Vector{T} #in pixels
end

function InstrumentalResponse()
    InstrumentalResponse(Float64[],Float64[],Float64[])
end

function (IR::InstrumentalResponse)(lambda,d=1.0)
    #d converts from pixels to wavelength
    sum([a*exp((-((d*l) - lambda)^2)/(d*w)^2) for (a,l,w) in zip(IR.amp,IR.location,IR.width)])
end

function kernel(lambda,IR::InstrumentalResponse,d=1.0)
    length(IR.location) == 0 && return Float64[]
    dL = lambda[2] - lambda[1]
    lmax = maximum(abs.(d*IR.location)) + 6*maximum(sqrt.(((d*IR.width).^2)/2))
    Lp = 0.0:dL:lmax
    L= vcat(-reverse(Lp)[1:end-1],Lp)
    k = IR.(L,d)
    return reflect(centered(k./sum(k))) #reflect for convolution not correlation
end

function kernel(lambda,IR::T,d=1.0) where T<:Real
    return reflect(Images.Kernel.gaussian((IR/d,), (21,)))
end

struct ExperimentalSpectra{T<:Real}
    lambda::Vector{T}
    data::Vector{T}
    err::Vector{T}
end

function Base.vcat(S::ExperimentalSpectra...)
    lambda = vcat((s.lambda for s in S)...)
    data = vcat((s.data for s in S)...)
    err = vcat((s.err for s in S)...)
    return ExperimentalSpectra(lambda,data,err)
end

struct TheoreticalSpectra{T<:Real}
    lambda::Vector{T}
    spec::Vector{T}
end

function Base.vcat(S::TheoreticalSpectra...)
    lambda = vcat((s.lambda for s in S)...)
    spec = vcat((s.spec for s in S)...)
    return TheoreticalSpectra(lambda,spec)
end


using OrthoMatchingPursuit
using LinearAlgebra
using Test

include("qr.jl")
include("elasticqr.jl")
include("pivotedqr.jl")
include("omp.jl")
include("wrapper.jl")


module SparseADRules

using ChainRulesCore: include
using Base: promote_eltype
using LinearAlgebra
using SparseArrays
using NiLang
using NiLang.AD
using ChainRulesCore

include("compat.jl")
include("linalg.jl")
include("lowranksvd.jl")
include("srcutils.jl")
include("chainrules.jl")
end


using Random, Statistics, ProgressMeter

#2x2-J1-J2-checkboard
function e_dif(model::Ising22, array, lattice, x, y)
    r = model.r
    a , b = mod(x,2), mod(y,2)

    top     = array[x, y + 1 - lattice * (y==lattice)]
    bottom  = array[x,y - 1 + lattice * (y==1)]
    left    = array[x - 1 + lattice * (x==1), y]
    right   = array[x + 1 - lattice * (x==lattice), y]

    if a==1
        right *= r
    else
        left *= r
    end
    if b==1
        top *= r
    else
        bottom *= r
    end

    return 2 * array[x, y] * (top+bottom+left+right)
end

#3x3-J1-J2
function e_dif(model::Ising33, array, lattice, x, y)
    r = model.r
    a , b = mod(x,3), mod(y,3)

    top     = array[x, y + 1 - lattice * (y==lattice)]
    bottom  = array[x,y - 1 + lattice * (y==1)]
    left    = array[x - 1 + lattice * (x==1), y]
    right   = array[x + 1 - lattice * (x==lattice), y]

    if a == 1
        if b == 2
            right *= r
        end
    elseif a == 2
        if b == 1
            top *= r
        elseif b == 2
            top *= r
            bottom *= r
            left *= r
            right *= r
        else
            bottom *= r
        end
    else
        if b == 2
            left *= r
        end
    end

    return 2 * array[x, y] * (top+bottom+left+right)
end

function onestep!(model::MT, spin_array, lattice, β) where {MT <: HamiltonianModel}
    for i = 1:lattice
        for j = 1:lattice
            e = e_dif(model, spin_array, lattice, i, j)
            if e <= 0
                spin_array[i,j] = -spin_array[i,j];
            elseif exp(-e*β) > rand()
                spin_array[i,j] = -spin_array[i,j];
            end
        end
    end
    return spin_array
end

function energy(model::MT, spin_array, lattice) where {MT <: HamiltonianModel}
    tol_en = 0
    for i = 1:lattice
        for j = 1:lattice
            e = e_dif(model, spin_array, lattice, i, j)
            tol_en += e
        end
    end
    return -tol_en/4
end

function MCMC(model::MT,lattice,β,Sweeps_heat,Sweeps) where {MT <: HamiltonianModel}
#     spin_array = ones(lattice,lattice)
    spin_array = (bitrand(lattice,lattice).-0.5)*2
    for j = 1:Sweeps_heat
        spin_array = onestep!(model,spin_array,lattice ,β);
    end
    mag = zeros(1,Sweeps);
    ene = zeros(1,Sweeps)
    Threads.@threads for j = 1:Sweeps
        spin_array = onestep!(model,spin_array,lattice ,β)
        mag[j] = abs(sum(spin_array)/lattice^2)
        ene[j] = energy(model,spin_array, lattice)/lattice^2
    end
    mag_ave = sum(mag)/Sweeps;
    ene_ave = sum(ene)/Sweeps;
    return mag_ave, ene_ave
end

function mutiMC(model::MT,β,lattice,bins,Sweeps_heat,Sweeps) where {MT <: HamiltonianModel}
    eMC_bins = zeros(bins,1)
    Threads.@threads for j=1:bins
        eMC_bins[j] =  MCMC(model,lattice,β,Sweeps_heat,Sweeps)
    end
    eMC = mean(eMC_bins)
    return eMC
end

function MutiMC(model::MT,β,lattice,Bins,bins,Sweeps_heat,Sweeps) where {MT <: HamiltonianModel}
    eMC_Bins = zeros(1,Bins)
    p = Progress(Bins)
    for j = 1:Bins
        eMC_Bins[j] = mutiMC(model,β,lattice,bins,Sweeps_heat,Sweeps)
        update!(p,j)
    end
    eMC_m = mean(eMC_Bins)
    err = 1.96*sqrt(var(eMC_Bins)/Bins)
    return eMC_m,err
end


using ITensorGaussianMPS
using LinearAlgebra
using ITensors

# Electrons

# Half filling
N = 100
Nf_up = N÷2
Nf_dn = N÷2
Nf = Nf_up + Nf_dn

@show N, Nf

# Maximum MPS link dimension
_maxlinkdim = 200

@show _maxlinkdim

# DMRG cutoff
_cutoff = 1e-8

# Hopping
t = 1.0

# Electron-electron on-site interaction
U = 1.0

@show t, U

# Make the free fermion Hamiltonian for the up spins
ampo_up = AutoMPO()
for n in 1:N-1
  ampo_up .+= -t, "Cdagup", n, "Cup", n+1
  ampo_up .+= -t, "Cdagup", n+1, "Cup", n
end

# Make the free fermion Hamiltonian for the down spins
ampo_dn = AutoMPO()
for n in 1:N-1
  ampo_dn .+= -t, "Cdagdn", n, "Cdn", n+1
  ampo_dn .+= -t, "Cdagdn", n+1, "Cdn", n
end

# Hopping Hamiltonians for the up and down spins
h_up = hopping_hamiltonian(ampo_up)
h_dn = hopping_hamiltonian(ampo_dn)

# Get the Slater determinant
Φ_up = slater_determinant_matrix(h_up, Nf_up)
Φ_dn = slater_determinant_matrix(h_dn, Nf_dn)

# Create an MPS from the slater determinants.
s = siteinds("Electron", N; conserve_qns = true)
println("Making free fermion starting MPS")
@time ψ0 = slater_determinant_to_mps(s, Φ_up, Φ_dn; eigval_cutoff = 1e-4,
                                                    cutoff = _cutoff,
                                                    maxdim = _maxlinkdim)
@show maxlinkdim(ψ0)

# The total non-interacting part of the Hamiltonian
ampo_noninteracting = AutoMPO()
for n in 1:N-1
  ampo_noninteracting .+= -t, "Cdagup", n, "Cup", n+1
  ampo_noninteracting .+= -t, "Cdagdn", n, "Cdn", n+1
  ampo_noninteracting .+= -t, "Cdagup", n+1, "Cup", n
  ampo_noninteracting .+= -t, "Cdagdn", n+1, "Cdn", n
end

H_noninteracting = MPO(ampo_noninteracting, s)
@show inner(ψ0, H_noninteracting, ψ0)
@show sum(diag(Φ_up' * h_up * Φ_up)) + sum(diag(Φ_dn' * h_dn * Φ_dn))

# The total interacting Hamiltonian
ampo_interacting = AutoMPO()
for n in 1:N-1
  ampo_interacting .+= -t, "Cdagup", n, "Cup", n+1
  ampo_interacting .+= -t, "Cdagdn", n, "Cdn", n+1
  ampo_interacting .+= -t, "Cdagup", n+1, "Cup", n
  ampo_interacting .+= -t, "Cdagdn", n+1, "Cdn", n
end
for n in 1:N
  ampo_interacting .+= U, "Nupdn", n
end
H = MPO(ampo_interacting, s)
#@show norm(prod(H) - prod(H_noninteracting))

# Random starting state
ψr = randomMPS(s, n -> n ≤ Nf ? (isodd(n) ? "↑" : "↓") : "0")

println("Random starting state energy")
@show flux(ψr)
@show inner(ψr, H, ψr)
println()
println("Free fermion starting state energy")
@show flux(ψ0)
@show inner(ψ0, H, ψ0)

println("\nStart from product state")
sweeps = Sweeps(10)
maxdim!(sweeps,10,20,_maxlinkdim)
cutoff!(sweeps,_cutoff)
er, ψ̃r = @time dmrg(H, ψr, sweeps)
@show er
@show flux(ψ̃r)

println("\nStart from free fermion state")
sweeps = Sweeps(5)
maxdim!(sweeps,_maxlinkdim)
cutoff!(sweeps,_cutoff)
e0, ψ̃0 = @time dmrg(H, ψ0, sweeps)
@show e0
@show flux(ψ̃0)

nothing


## Model definition

""" `Model` objects assume that there is a pure-kwarg constructors (such as provided by
Parameters.jl) """
abstract type Model end

# Models should implement these functions (or rely on the defaults)
transition_mat(m::Model, inputs, i) = Identity()
transition_noise(m::Model, inputs, i) = no_noise()
observation_mat(m::Model, inputs, i) = Identity()
function observation_noise end  # necessary for all models
function observation end        # optional; you can pass `observations=...` instead
function initial_state end      # optional; you can pass `initial_state=...` instead

# More exotic model functions
skip_observation(m::Model, inputs, i) = false
nobservations(m::Model, inputs) = size(observation_mat(m, inputs, 1), 1) #fails if skipped
function labels end             # for pretty-printing
parameters(m::M) where M<:Model = fieldnames(M)

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

predicted_state(state_prior::Gaussian, transition_mat, transition_noise::Gaussian) =
    Gaussian(transition_mat * mean(state_prior) + mean(transition_noise),
             transition_mat * cov(state_prior) * transition_mat' + cov(transition_noise))

predicted_state(prev_state::Gaussian, m::MiniKalman.Model, inputs, t::Int) =
    predicted_state(prev_state, transition_mat(m, inputs, t),
                    transition_noise(m, inputs, t))

function predicted_observation(state::Gaussian, observation_mat, observation_noise)
    # Prediction step (μ_(t|t-1), Σ_(t|t-1))
    C = observation_mat
    Du, R = observation_noise     # Du := Dₜuₜ
    μ, Σ = state
    ŷ = C * μ + Du               # Murphy forgot the Du term in his book
    S = C * Σ * C' + R
    return Gaussian(ŷ, S)
end

predicted_observation(state::Gaussian, m::MiniKalman.Model, inputs, t::Int) =
    predicted_observation(state, observation_mat(m, inputs, t),
                          observation_noise(m, inputs, t))

function kfilter_obs(state::Gaussian,
                     observation, observation_mat, observation_noise::Gaussian)
    # Following Machine Learning, A probabilistic perspective, by Kevin Murphy, 18.3.1.2,
    # with the exception that I'm setting the forcing term to 0, but allowing noise terms
    # with means (without loss of generality)

    C = observation_mat
    y = observation
    μ, Σ = state

    ŷ, S = predicted_obs =
        predicted_observation(state, observation_mat, observation_noise)

    # Filter
    K = Σ * C' / S         # Kalman gain matrix
    r = y - ŷ
    filtered_state = Gaussian(μ + K*r,
                              (I - K*C) * Σ)
    ll = logpdf(predicted_obs, y)  # log-likelihood
    return (filtered_state, ll, predicted_obs)
end

""" Perform one step of Kalman filtering, for online use. We assume equations:

```julia
    current_state = transition_mat * previous_state + transition_noise
    observation = observation_mat * current_state + observation_noise
```

and return this tuple:

```
   (current_state::Gaussian,   # the posterior `P(state|observation)`
    ll,                        # the log-likelihood for that step
    predicted_obs::Gaussian)   # `P(observation|state)`
```

To add an "input" term, pass it as `transition_noise = Gaussian(input_term, noise_cov)`.
"""
function kfilter(prev_state::Gaussian, m::MiniKalman.Model, inputs, t::Int,
                 observations=nothing)
    state = predicted_state(prev_state, m, inputs, t)
    if skip_observation(m, inputs, t)
        # Optimization note: returning ll=0.0 might make this function type-unstable when
        # computing gradients, for gradient descent.
        # We can't have `kfilter_obs` run when `skip_observation` is true, because
        # bad matrices (singular / containing missing) can trigger exceptions.
        (state, 0.0, predicted_observation(state, m, inputs, t))
    else
        kfilter_obs(state,
                    observations===nothing ? observation(inputs, t) : observations[t],
                    observation_mat(m, inputs, t), observation_noise(m, inputs, t))
    end
end

function prediction_type(m::Model, inputs, any_state::Gaussian)
    # This juggling is necessary to get the proper `SVector/Vector` output
    # type. It's lame, but I don't see how else to do it (except maybe using `accumulate`,
    # but IIRC there were some performance concerns)
    fake_obs_mat = hcat((mean(any_state) for _ in 1:nobservations(m, inputs))...)'
    fake_obs = fake_obs_mat * mean(any_state)
    return Gaussian{typeof(fake_obs), typeof(fake_obs*fake_obs')}
end

""" (Internal function) return three vectors, appropriate for storing
states, likelihoods, and P(observation). """
function output_vectors(m::Model, inputs, observations=nothing; length=length(inputs))
    # Note: this function was split off for our own internal purposes (switching states)
    # It could be merged back into `kalman_filter!`, I suppose.
    state = make_full(initial_state(m))
    filtered_states = Vector{typeof(state)}(undef, length)
    predicted_obs = Vector{prediction_type(m, inputs, state)}(undef, length)
    lls = Vector{Float64}(undef, length)
    return (filtered_states, lls, predicted_obs)
end

""" Perform Kalman filtering and store the results in the `filtered_states, lls,
predicted_obs` vectors, for the given `steps`. """
function kalman_filter!(filtered_states::AbstractVector, lls::AbstractVector,
                        predicted_obs::AbstractVector,
                        m::Model, inputs, observations=nothing;
                        steps::AbstractRange=1:length(filtered_states),
                        initial_state=(steps[1]==1 ? initial_state(m) :
                                       filtered_states[steps[1]-1]))
    state = make_full(initial_state)  # we need make_full so that the state does
                                      # not change type during iteration
    for t in steps
        state, lls[t], predicted_obs[t] = kfilter(state, m, inputs, t, observations)
        filtered_states[t] = state
    end
end

function kalman_filter(m::Model, inputs, observations=nothing;
                       initial_state=initial_state(m),
                       steps=1:length(observations===nothing ? inputs : observations))
    out_vecs = output_vectors(m, inputs, observations, length=length(steps))
    kalman_filter!(out_vecs..., m, inputs, observations; steps=steps,
                   initial_state=initial_state)
    return out_vecs
end

function log_likelihood(m::Model, inputs, observations=nothing;
                        initial_state=MiniKalman.initial_state(m),
                        steps=1:length(observations===nothing ? inputs : observations))
    # Since this is in the inner loop of `optimize`, we make sure it's non-allocating
    ll_sum = 0.0
    state = make_full(initial_state)
    for t in steps
        state, ll, _ = kfilter(state, m, inputs, t, observations)
        ll_sum += ll
    end
    return ll_sum
end

kalman_filtered(args...; kwargs...) = kalman_filter(args...; kwargs...)[1]  # convenience
kalman_predicted(args...; kwargs...) = kalman_filter(args...; kwargs...)[3]  # convenience

""" Convenience; returns a vector of the total likelihood up to each step. """
cumulative_log_likelihood(args...; kwargs...) =
    cumsum(kalman_filter(args...; kwargs...)[2])


#= The Computer Language Benchmarks Game
   https://salsa.debian.org/benchmarksgame-team/benchmarksgame/

   contributed by Ben Lauwens

   derived from the Chapel version by Tom Hildebrandt, Brad Chamberlain, and Lydia Duncan
   and inspired by the Julia version of Luiz M. Faria
=#

using ResumableFunctions

using Base.GMP.MPZ: add_ui!, mul_ui!, add!, tdiv_q!
const mpz_t = Ref{BigInt}
addmul_ui!(x::BigInt,a::BigInt,b::UInt64) = ccall((:__gmpz_addmul_ui,"libgmp"), Cvoid, (mpz_t,mpz_t,Culong), x, a, b)
submul_ui!(x::BigInt,a::BigInt,b::UInt64) = ccall((:__gmpz_submul_ui,"libgmp"), Cvoid, (mpz_t,mpz_t,Culong), x, a, b)
get_ui(x::BigInt) = ccall((:__gmpz_get_ui,"libgmp"), Culong, (mpz_t,), x)

@resumable function pidigits()
    k = one(UInt64)
    d = zero(UInt64)
    numer = BigInt(1)
    denom = BigInt(1)
    accum = BigInt(0)
    tmp1 = BigInt(0)
    tmp2 = BigInt(0)
    tmp3 = BigInt(0)
    while true
        k2 = 2k + one(UInt64)
        addmul_ui!(accum, numer, UInt64(2))
        mul_ui!(accum, k2)
        mul_ui!(denom, k2)
        mul_ui!(numer, k)
        k += one(UInt64)
        if numer <= accum
            mul_ui!(tmp1, numer, UInt64(3))
            add!(tmp1, accum)
            tdiv_q!(tmp2, tmp1, denom)
            mul_ui!(tmp1, numer, UInt64(4))
            add!(tmp1, accum)
            tdiv_q!(tmp3, tmp1, denom)
            if tmp2 == tmp3
                d = get_ui(tmp2)
                @yield d
                submul_ui!(accum, denom, d)
                mul_ui!(accum, UInt64(10))
                mul_ui!(numer, UInt64(10))
            end
        end
    end
end

function main(n = parse(Int64, ARGS[1]))
    i = 0
    pid = pidigits()
    for i in 1:n
        d = pid()
        print(d)
        if i % 10 === 0
            println("\t:", i)
        end
    end
    if i % 10 !== 0
        println(" "^(10 - (i % 10)), "\t:", n)
    end
end

main(100)


using Mimi
using OptiMimi
using Test

# Easier basic model to optimize:
# Each period, you have 1 unit to divide between consuming and staving off disaster.
# If disaster happens, you get 0 consumption for that period.
# Utility is defined as sqrt(c).  The probability of disaster is 1-c in each period.
# Assume 5% discounting and 10 periods.

## Analytical solution in 1 period

## max_c (1 - c) u(c) => c = 1/3
gamma = exp(-.05) # 5% discount rate

# OptiMimi solution

# Create simple model
@defcomp Bellmano begin
    # The x-value to evaluate the quadratic
    consumption = Parameter(index=[time])

    # The y-value of the quadratic at the x-value
    utility = Variable(index=[time])

    function run_timestep(p, v, d, t)
        #v.utility[tt] = sqrt(p.consumption[tt]) * (1 - p.consumption[tt])

        if rand() < p.consumption[tt]
            v.utility[tt] = 0
        else
            v.utility[tt] = sqrt(max(0, p.consumption[tt]))
        end
    end
end

m = Model()
set_dimension!(m, :time, collect(1:10))

bellmano = add_comp!(m, Bellmano)
bellmano[:consumption] = repmat([.25], 10)

run(m)

import OptiMimi.uncertainproblem
using DataFrames

df = DataFrame(algorithm=[], mcperlife=[], cons1=[], cons2=[], cons3=[], cons4=[], cons5=[], cons6=[], cons7=[], cons8=[], cons9=[], cons10=[], cons1serr=[], cons2serr=[], cons3serr=[], cons4serr=[], cons5serr=[], cons6serr=[], cons7serr=[], cons8serr=[], cons9serr=[], cons10serr=[], time=[])
push!(df, [:exact; Inf; repmat([1/3], 10); repmat([0], 10); 0])

for samplemult in [1, 2, 4]
    for mcperlife in [5, 20, 50]
        println(mcperlife)
        prob = uncertainproblem(m, [:Bellmano], [:consumption], [0.], [1.], m -> sum(m[:Bellmano, :utility] .* exp(-(0:9) * .05)), (model) -> nothing)

        tic()
        soln = solution(prob, () -> repmat([.25], 10), :social, mcperlife, samplemult)
        time = toc()
        push!(df, [:social; mcperlife; soln.xmean; soln.xserr; time])

        println(df)

        tic()
        soln = solution(prob, () -> repmat([.25], 10), :biological, mcperlife, 120000*samplemult)
        time = toc()
        push!(df, [:biological; mcperlife; soln.xmean; soln.xserr; time])

        println(df)

        tic()
        soln = solution(prob, () -> repmat([.25], 10), :sampled, mcperlife, 300*samplemult)
        time = toc()
        push!(df, [:sampled; mcperlife; soln.xmean; soln.xserr; time])

        println(df)
    end

    writetable("attempts-compare2.csv", df)
end


include("basic_setup.jl")

function perf_est(subsample)
    algorithm = AnnealingEstimate(subsample, 50, 100)
    Trajectories.simulate(algorithm, initial, system)
end

results = perf_est.(1:5:26)

using Statistics
using Plots

plot(1:5:26, Trajectories.log_marginal.(results), yerr=sqrt.(var.(results)))
plot(1:5:26, sqrt.(var.(results)))

plot(result.inv_temps[1:end - 1], mean(result.acceptance, dims=2))
plot(result.inv_temps[begin + 1:end - 1], diff(mean(result.weights, dims=2), dims=1))

histogram(result.weights[end, :])

Trajectories.log_marginal(result)
sqrt(var(result))

mean(result.weights[end, :])

var(result.weights[end, :]) / 10_000


module CoarseGrainingPrism3D
#####
using LinearAlgebra
using Profile
#####
#####

### GLOBAL CONSTANT
const K =3
const x = K+1
const y = K/2

### FUNDAMENTAL FUNCTION
#chop function
export numchop
function numchop(x::Real)
    if abs(x) <= 1000*eps()
        x = 0
    end
    return x
end
function numchop(x::ComplexF64)
    ansreal = real(x)
    ansim = imag(x)
    if abs(ansreal) <= 1000*eps()
        ansreal = 0
    end
    if abs(ansim) <= 1000*eps()
        ansim = 0
    end
    return x= ansreal + ansim*im
end
function numchop(x::Array)
    for i in 1:length(x)
        if x[i] != 0
            x[i] = numchop(x[i])
        end
    end
    return x
end

#Define delta_{ijkl} -> coupling rules
export delta
function delta(i::Real,j::Real,k::Real)
	sol = 0
	if i <= (j+k) && j<= (i+k) && i+j+k <= K &&2*(i+j+k)%2 == 0
		sol =1
	end
	return sol
end

#Define quantum numbers qn (this is real)
export qn
function qn(n::Real)
    sol = (exp(pi*n*im/(K+2)) - exp(-pi*n*im/(K+2))) / (exp(pi*im/(K+2)) - exp(-pi*im/(K+2)))
    return real(sol)
end

#Define qn factorial
export qnfact
function qnfact(n::Real)
    sol = 1
    for i in 1:n
        sol *= qn(i)
    end
    return sol
end

#Define square root of quantum dimension
export visqrt
function visqrt(i::Real)
    sol = ((-1+0im)^i )*sqrt(qn(2*i+1))
    return sol
end

#Define triangle equality
export trian
function trian(i::Real,j::Real,k::Real)
    sol = 0
    if delta(i,j,k) == 1
        sol = delta(i,j,k)*sqrt(qnfact(i+j-k)*qnfact(i-j+k)*qnfact(-i+j+k)/qnfact(i+j+k+1))
    end
    return sol
end

# Define Racah-Wigner six-j symbol
export RacWig6j
function RacWig6j(i::Real,j::Real,m::Real,k::Real,l::Real,n::Real)
    a = i+j+m; b = i+l+n; c = k+j+n; d = k+l+m;  e = i+j+k+l; f = i+k+m+n; g = j+l+m+n
    sol = 0
    if delta(i,j,m) != 0 && delta(i,l,n) != 0 && delta(k,j,n) != 0 && delta(k,l,m) != 0
        sumz = 0
        for z in max(a,b,c,d):min(e,f,g)
            sumz += (-1)^z *qnfact(z+1)/
                ((qnfact(e-z)*qnfact(f-z)*qnfact(g-z))* (qnfact(z-a)*qnfact(z-b)*qnfact(z-c)*qnfact(z-d)))
        end
        sol = trian(i,j,m)*trian(i,l,n)*trian(k,j,n)*trian(k,l,m)*sumz
    end
    return sol
end

#Define F-symbol
export Fsymb
function Fsymb(i::Real,j::Real,m::Real,k::Real,l::Real,n::Real)
    sol = 0
    if delta(i,j,m) != 0 && delta(i,l,n) != 0 && delta(k,j,n) != 0 && delta(k,l,m) != 0
        sol = (-1+0im)^(i+j+k+l)*sqrt(qn(2*m+1)*qn(2*n+1)) * RacWig6j(i,j,m,k,l,n)
    end
    return sol
end

#Define G-symbol
export Gsymb
function Gsymb(i::Real,j::Real,m::Real,k::Real,l::Real,n::Real)
    sol = 0
    if delta(i,j,m) != 0 && delta(i,l,n) != 0 && delta(k,j,n) != 0 && delta(k,l,m) != 0
        sol = Fsymb(i,j,m,k,l,n) /(visqrt(m)*visqrt(n))
    end
    return sol
end

### DATA
# Define amplitude for a prism function of 12 edges : usual 9 + 3 diag

export Prism
function Prism(β::Float64,α1::Float64,β1::Float64,α2::Float64,δ::Float64,a1::Float64,b::Float64,
                    a2::Float64,d::Float64,c1::Float64,c2::Float64,f::Float64)
    sol = 0
    if delta(c2,d,δ) !=0 && delta(c2,α2,a2) !=0 && delta(β,d,a2) !=0 && delta(β,α2,δ) !=0 &&
        delta(f,d,α1) !=0 && delta(f,c1,b) !=0 && delta(a1,d,b) !=0 && delta(a1,c1,α1) !=0 &&
        delta(α1,β1,δ) !=0 && delta(c2,β1,f) !=0

        dims = visqrt(β)*visqrt(α1)*visqrt(β1)*visqrt(α2)*visqrt(δ)*visqrt(a1)*visqrt(b)*visqrt(a2)*visqrt(d)*visqrt(c1)*visqrt(c2)*visqrt(f)
        sol = dims*Gsymb(c2,d,δ,β,α2,a2)*Gsymb(f,d,α1,a1,c1,b)*Gsymb(c2,d,δ,α1,β1,f)

    if imag(sol) == 0
        sol = real(sol)
    end
    end
    return sol
end

export SuperIndex
function SuperIndex(AllSpin)
	Size = length(AllSpin)
	ans = Int(2*sum(AllSpin[i]*x^(Size-i) for i in 1:Size)+1)
    return ans
end

export InverseSuperIndex
function InverseSuperIndex(y,Size)
    yt = y - 1
    j = zeros(0)
    for i in (Size-1):-1:0
        d,r = divrem(yt,x^i)
        push!(j,d/2)
        yt = r
    end
    return j
end

export PrAOpti
function PrAOpti()
    PrA = Array{Float64}(undef,0)
	s = Array{Int64}(undef,0)

	sol = 0
	dims = 0
    for β in 0.:0.5:y, d in 0.:0.5:y, a2 in 0.:0.5:y
        if delta(β,d,a2) !=0
            for α2 in 0.:0.5:y, δ in 0.:0.5:y
                if  delta(β,α2,δ) !=0
                    for c2 in 0.:0.5:y
                        if delta(c2,d,δ) !=0 && delta(c2,α2,a2) !=0
                            for f in 0.:0.5:y, β1 in 0.:0.5:y
                                if delta(c2,β1,f) !=0
                                    for α1 in 0.:0.5:y
                                        if delta(f,d,α1) !=0 && delta(α1,β1,δ) !=0
                                            for c1 in 0.:0.5:y, a1 in 0.:0.5:y
                                                if delta(a1,c1,α1) !=0
                                                    for b in 0.:0.5:y
                                                        if delta(f,c1,b) !=0 && delta(a1,d,b) !=0
                                                            dims = visqrt(β)*visqrt(α1)*visqrt(β1)*visqrt(α2)*visqrt(δ)*visqrt(a1)*visqrt(b)*visqrt(a2)*visqrt(d)*visqrt(c1)*visqrt(c2)*visqrt(f)
                                                            sol = numchop(dims*Gsymb(c2,d,δ,β,α2,a2)*Gsymb(f,d,α1,a1,c1,b)*Gsymb(c2,d,δ,α1,β1,f))
                                                            if sol !=0
                                                                push!(PrA,sol)
                                                                push!(s,SuperIndex([β,α1,β1,α2,δ,a1,b,a2,d,c1,c2,f]))
                                                            end
                                                        end
                                                    end
                                                end
                                            end
                                        end
                                    end
                                end
                            end
                        end
                    end
                end
            end
        end
    end
	return PrA, s
end

##### coarse-graining algorithm
# create function returning M_{AB}^C, where C fixed spins, A spins of tetra and B spins of pyramid
export BlocksPrism
function BlocksPrism(PrismData,c2::Float64,d::Float64,δ::Float64)
    PrA = PrismData[1]
    sIndex = PrismData[2]

    TempTetra = Array{Array}(undef,0)
    TempPyra = Array{Array}(undef,0)
    TempAmp = Array{Real}(undef,0)
    PosAmp = 1
    for i in sIndex
        (β,α1,β1,α2,δCut,a1,b,a2,dCut,c1,c2Cut,f) = InverseSuperIndex(i,12)
        if δCut == δ && dCut == d && c2Cut == c2
            push!(TempAmp,PrA[PosAmp])
            push!(TempTetra,[β α2 a2])
            push!(TempPyra,[c1 a1 α1 β1 b f])
        end
        PosAmp += 1
    end
    PosTetra = unique(TempTetra)
    PosPyra = unique(TempPyra)
    SizeA = length(PosTetra)
    SizeB = length(PosPyra)
    mat = zeros(SizeA,SizeB)
    for i in 1:length(TempTetra)
        qa = findall(x -> x == TempTetra[i], PosTetra)
        qb = findall(x -> x == TempPyra[i], PosPyra)
        mat[qa[1],qb[1]] = TempAmp[i]
    end

    if SizeA == 0
        mat = 0
    end
    return mat, PosTetra, PosPyra
end

export svdPrism
function svdPrism(PrismData,c2::Float64,d::Float64,δ::Float64)# parameters c2,d,δ
    mat, TetraIndex, PyraIndex = BlocksPrism(PrismData,c2,d,δ)

    U, s, V = svd(mat)
    V = adjoint(V)

    return U, numchop(s), V, TetraIndex, PyraIndex
end

#HYP : only take the firsts TRUNC singular values
export svdPrismTruncated
function svdPrismTruncated(PrismData,c2::Float64,d::Float64,δ::Float64,trunc)
	U, s, V, TetraIndex, PyraIndex = svdPrism(PrismData,c2,d,δ)

	TruncU = U[:,1:trunc]*sqrt.(s[1:trunc])
	TruncV = transpose(sqrt.(s[1:trunc]))*V[1:trunc,:]

	if real(sqrt(visqrt(c2)*visqrt(d)*visqrt(δ))) == 0
		TruncU = im*sqrt(visqrt(c2)*visqrt(d)*visqrt(δ))*TruncU
		TruncV = -im*sqrt(visqrt(c2)*visqrt(d)*visqrt(δ))*TruncV
	else
		TruncU = sqrt(visqrt(c2)*visqrt(d)*visqrt(δ))*TruncU
		TruncV = sqrt(visqrt(c2)*visqrt(d)*visqrt(δ))*TruncV
	end
    if length(TetraIndex) != 0
        (β,α2,a2) = TetraIndex[1]
        if sign(TruncU[1,1]) != sign(AmpTetraInPrism(c2,d,δ,β,α2,a2))
             TruncU = -(1)*TruncU
             TruncV = -(1)*TruncV
        end
    end

	return real(TruncU), real(TruncV), TetraIndex, PyraIndex, s[1:trunc]
end

export FullsvdPrismTruncated
function FullsvdPrismTruncated(PrismData,trunc)
	FullU = Array{Array}(undef,0)
	FullV = Array{Array}(undef,0)
	FullTetraindex = Array{Array}(undef,0)
	FullPyraindex = Array{Array}(undef,0)
	FullCutindex = Array{Array}(undef,0)
	Fulls = Array{Array}(undef,0)
	for c2 in 0:0.5:y, d in 0:0.5:y, δ in 0:0.5:y
		if delta(c2,d,δ) == 1
			temp = svdPrismTruncated(PrismData,c2,d,δ,trunc)
            if temp[1] !=[0]
    			push!(FullU,temp[1])
    			push!(FullV,temp[2])
    			push!(FullTetraindex,temp[3])
    			push!(FullPyraindex,temp[4])
    			push!(FullCutindex,[c2,d,δ])
    			push!(Fulls,temp[5])
            end
		end
	end

    if FullU[1][1] != 1
        FullU = (1/FullU[1][1])*FullU
        FullV = (1/FullV[1][1])*FullV
    end

	return FullU, FullV, FullTetraindex, FullPyraindex, FullCutindex, Fulls
end

export PrismUVTruncated
function PrismUVTruncated(PrismData,trunc)
	FullU, FullV, FullTetraindex, FullPyraindex, FullCutindex,s = FullsvdPrismTruncated(PrismData,trunc)

	AmpUV = Array{Float64}(undef,0)
	PosUVFaceExt = Array{Int64}(undef,0)
	PosUVFaceTrue = Array{Int64}(undef,0)
	PosUVFaceFalse = Array{Int64}(undef,0)
	SizeC = length(FullCutindex)
	for i in 1:SizeC, j in 1:SizeC
		(c2,d,δ) = FullCutindex[i]
		(c2t,dt,δt) = FullCutindex[j]
		#glue along dA/dB
		if d == dt
			TempTetra = FullTetraindex[i]
			TempPyra = FullPyraindex[j]
			for iA in 1:length(TempTetra), jB in 1:length(TempPyra)
				(β,α2,a2) = TempTetra[iA]
				(c1t,a1t,α1t,β1t,bt,ft) = TempPyra[jB]
				#and glue along uA/uB and vA/vB
				if β == bt && a2 == a1t
					push!(AmpUV,FullU[i][iA]*FullV[j][jB]/(visqrt(d)*visqrt(β)*visqrt(a2)))
					push!(PosUVFaceExt,SuperIndex([β1t,ft,β]))
					push!(PosUVFaceTrue,SuperIndex([c2,δt,c2t,δ,d]))
					push!(PosUVFaceFalse,SuperIndex([c2,α1t,c1t,α2,a2]))
				end
			end
		end
	end
	return AmpUV,PosUVFaceExt,PosUVFaceTrue,PosUVFaceFalse, s
end

export PrismEff
function PrismEff(PrismData,trunc)

	AmpUV,PosUVFaceExt,PosUVFaceTrue,PosUVFaceFalse, s = PrismUVTruncated(PrismData,trunc)
	FullAmp = Array{Real}(undef,0)
	BigSIndex = Array{Real}(undef,0)
	for i in 1:length(PosUVFaceTrue)
		# spins of first prism
		(Χ,c1,a2,c2,d) = InverseSuperIndex(PosUVFaceTrue[i],5) #shared spins
		(γ,e,f) = InverseSuperIndex(PosUVFaceExt[i],3)
		(χ,α1,a1,α2,β) = InverseSuperIndex(PosUVFaceFalse[i],5) # was face 2

		pos = findall(x->x == PosUVFaceTrue[i],PosUVFaceFalse)
		for j in pos
			# spins of second prism
			#(bA,j1A,j4B,j3B,uA)
			(γt,et,ft) = InverseSuperIndex(PosUVFaceExt[j],3)
			(Χ,c1t,a2t,c2t,dt) = InverseSuperIndex(PosUVFaceTrue[j],5)
			#super index for 18 spins prism
			TempIndex = SuperIndex([χ,α1,a1,α2,β,c1t,a2t,c2t,dt,γ,e,f,γt,et,ft,c1,c2,a2])
			# FullTot = vcat(FullTot,[TempIndex AmpUV[i]*AmpUV[j]])
			push!(FullAmp,AmpUV[i]*AmpUV[j])
			push!(BigSIndex,TempIndex)
		end
	end

    p = sortperm(BigSIndex)
    BigSIndex = BigSIndex[p]
    FullAmp = FullAmp[p]

    NewSindex = Array{Int64}(undef,0)
    NewAmp = Array{Float64}(undef,0)

    cnt = 0
    for i in 2:length(BigSIndex)
        cnt +=1
        if BigSIndex[i] == BigSIndex[i-1] && i != length(BigSIndex)
            FullAmp[i-cnt] += FullAmp[i]
        else
            push!(NewAmp,FullAmp[i-cnt])
            push!(NewSindex,BigSIndex[i-1])
            cnt = 0
        end
    end
    if BigSIndex[end] == BigSIndex[end-1]
        NewAmp[end] += FullAmp[end]
    else
        push!(NewAmp,FullAmp[end])
        push!(NewSindex,BigSIndex[end])
    end
    return NewAmp, NewSindex
end

####
"embedding map"
####

# 2-2 move for diag
# WARNING: Move done with Fsym, wich need to be either redefine at every step ? or always keep
# initial one ? but then we know something weird might happen -> do test !! IN THE FOLLOWING
# DONE WITH INITIAL Fsymb

export Move2_2
function Move2_2(SetAmpAndIndex,SizeSetSpins,PosMove,PosTrianglesOfMove)
	# SetIndexAndAmp: full set of amplitude and SuperIndex pos of non zero
	# SizeSetSpins: number of edges in polyhedra associated to SetIndexAndAmp
	# PosMove: Position of the edge/spin that we want to flip
	# PosTrianglesOfMove: Positions of the OTHER 4 edges, forming with PosMove the two ini triangles
	# 				that is PosTrianglesOfMove = [j1,j2 , h1,h2] and j1j2Posmove / h1h2Posmove ini triangles
	#				AND j1,h1 / j2,h2 share a vertex
	Amp,sIndex = SetAmpAndIndex

	TempAmp = Array{Real}(undef,0)
	TempsIndex = Array{Real}(undef,0)
	for i in 1:length(sIndex)
		SetSpins = InverseSuperIndex(sIndex[i],SizeSetSpins)
		for qMove in 0:0.5:y
            TempFsymb = Fsymb(SetSpins[PosTrianglesOfMove[1]],SetSpins[PosTrianglesOfMove[2]],SetSpins[PosMove],SetSpins[PosTrianglesOfMove[4]],SetSpins[PosTrianglesOfMove[3]],qMove)
            if numchop(TempFsymb) != 0
                TempSetSpins = copy(SetSpins)
				push!(TempsIndex,SuperIndex(insert!(deleteat!(TempSetSpins,PosMove),PosMove,qMove)))
				push!(TempAmp,TempFsymb*Amp[i])
			end
		end
	end

    p = sortperm(TempsIndex)
    TempsIndex = TempsIndex[p]
    TempAmp = TempAmp[p]

    NewsIndex = Array{Int64}(undef,0)
	NewData = Array{Float64}(undef,0)

	cnt = 0
	for i in 2:length(TempsIndex)
        cnt +=1
		if TempsIndex[i] == TempsIndex[i-1] && i != length(TempsIndex)
			TempAmp[i-cnt] += TempAmp[i]
		else
            if numchop(TempAmp[i-cnt]) != 0
                push!(NewData,TempAmp[i-cnt])
                push!(NewsIndex,TempsIndex[i-1])
            end
		    cnt = 0
		end
	end
    if TempsIndex[end] == TempsIndex[end-1]
        NewData[end] += TempAmp[end]
    else
        push!(NewData,TempAmp[end])
        push!(NewsIndex,TempsIndex[end])
    end

	return NewData,NewsIndex
end

export Move3_1
function Move3_1(SetAmpAndIndex,SizeSetSpins,PosTrianglesOfMove)
    # SetIndexAndAmp: full set of amplitude and SuperIndex pos of non zero
	# SizeSetSpins: number of edges in polyhedra associated to SetIndexAndAmp
	# PosTrianglesOfMove: Position of the 6 edges of the move, of the form
    #                   [j1 j2 j3 h1 h2 h3] where ji are disappearing spins, h1 h2 h3 keeping spins
    #                                triangles are: j1j3h1, j1j2h2, j2j3h3
	Amp,sIndex = SetAmpAndIndex

    PositionDelete = Array{Bool}(undef,0)
    for i in 1:SizeSetSpins
        if i in PosTrianglesOfMove[1:3]
            push!(PositionDelete,true)
        else
            push!(PositionDelete,false)
        end
    end

    TempAmp = Array{Real}(undef,0)
	TempsIndex = Array{Real}(undef,0)
    for i in 1:length(sIndex)
        SetSpins = InverseSuperIndex(sIndex[i],SizeSetSpins)
        TempDimFactor = (visqrt(SetSpins[PosTrianglesOfMove[1]])*visqrt(SetSpins[PosTrianglesOfMove[3]]))/visqrt(SetSpins[PosTrianglesOfMove[4]])
        TempFsymb =Fsymb(SetSpins[PosTrianglesOfMove[6]],SetSpins[PosTrianglesOfMove[5]],SetSpins[PosTrianglesOfMove[4]],SetSpins[PosTrianglesOfMove[1]],SetSpins[PosTrianglesOfMove[3]],SetSpins[PosTrianglesOfMove[2]])
        if numchop(TempFsymb) != 0
            TempSetSpins = copy(SetSpins)
            push!(TempsIndex,SuperIndex(deleteat!(TempSetSpins,PositionDelete)))
            push!(TempAmp,(1/(TempDimFactor*TempFsymb))*Amp[i])
        end
    end

   p = sortperm(TempsIndex)
   TempsIndex = TempsIndex[p]
   TempAmp = TempAmp[p]

   NewsIndex = Array{Int64}(undef,0)
   NewData = Array{Real}(undef,0)
   cntlast = 0
   cnt = 0
   for i in 2:length(TempsIndex)
       cnt+=1
       if TempsIndex[i] == TempsIndex[i-1] && i != length(TempsIndex)
           #TempAmp[i-cnt] += TempAmp[i]
       else
           if numchop(TempAmp[i-cnt]) != 0
               #push!(NewData,(1/cnt)*TempAmp[i-cnt])
               push!(NewData,TempAmp[i-cnt])
               push!(NewsIndex,TempsIndex[i-1])
               if i == length(TempsIndex)
                   cntlast=cnt
               end
           end
           cnt = 0
        end
    end
    if TempsIndex[end] == TempsIndex[end-1]
        #NewData[end] = cntlast*NewData[end]
        #NewData[end] += TempAmp[end]
        #NewData[end] = (1/(cntlast+1))*NewData[end]
    else
        push!(NewData,TempAmp[end])
        push!(NewsIndex,TempsIndex[end])
    end

    if NewData[1] != 1
        NewData = (1/NewData[1])*NewData
    end

	return NewData,NewsIndex
end

#### Embedding of PrismEff to space of Prism again
export EmbPrismEff
function EmbPrismEff(PrismData,trunc)
    SetAmpAndIndex = PrismEff(PrismData,trunc)

    #Move2_2 for big diagonal
    #Move2_2_BigDiag = Move2_2(SetAmpAndIndex,18,18,[10,11,14,15])
    Move2_2FirstTriangle = Move2_2(SetAmpAndIndex,18,16,[2,10,6,13])

    # Move 2_2 first triangle
    #Move2_2FirstTriangle = Move2_2(Move2_2_BigDiag,18,16,[2,10,6,13])
    Move2_2_BigDiag = Move2_2(Move2_2FirstTriangle,18,18,[10,11,14,15])

    # Move3_1 first triangle
    #Move3_1FirstTriangle = Move3_1(Move2_2FirstTriangle,18,[10,13,14,18,16,7])
    Move3_1FirstTriangle = Move3_1(Move2_2_BigDiag,18,[10,13,14,18,16,7])

    # Move 2_2 second triangle: SHOULD OR NOT ?????
    Move2_2SecondTriangle = Move2_2(Move3_1FirstTriangle,15,14,[4,11,8,12])

    # Move3_1 second triangle: SHOULD OR NOT ?????
    Move3_1SecondTriangle = Move3_1(Move2_2SecondTriangle,15,[10,11,12,15,3,14])

    EffPrism = Move3_1SecondTriangle

    return EffPrism
end

##### test function

#### data from initial tetra in prism
export AmpTetraInPrism
function AmpTetraInPrism(j4::Real,j5::Real,j6::Real,j1::Real,j2::Real,j3::Real)
	ans = 0
	if delta(j4,j5,j6)==1 && delta(j4,j2,j3)==1 && delta(j1,j5,j3)==1 && delta(j1,j2,j6)==1
		dims = visqrt(j1)*visqrt(j2)*visqrt(j3)*visqrt(j4)*visqrt(j5)*visqrt(j6)
		ans = dims*Gsymb(j4,j5,j6,j1,j2,j3)
	end
	return ans
end

export FullAmpTetraInPrism
function FullAmpTetraInPrism()
    Amp = Array{Real}(undef,0)
    Index = Array{Real}(undef,0)
    for j1 in 0:0.5:y, j2 in 0:0.5:y, j3 in 0:0.5:y, j4 in 0:0.5:y, j5 in 0:0.5:y, j6 in 0:0.5:y
        TempAmp = numchop(AmpTetraInPrism(j1,j2,j3,j4,j5,j6))
        if TempAmp!=0
            push!(Amp,TempAmp)
            push!(Index,SuperIndex([j1,j2,j3,j4,j5,j6]))
        end
    end
    return Amp,Index
end

export FullAmpPyramidInPrism
function FullAmpPyramidInPrism()
	Amp = Array{Real}(undef,0)
	Index = Array{Real}(undef,0)
	for j1 in 0:0.5:y, j2 in 0:0.5:y, j3 in 0:0.5:y, j4 in 0:0.5:y, j5 in 0:0.5:y, j6 in 0:0.5:y,
		a in 0:0.5:y, b in 0:0.5:y, c in 0:0.5:y
		TempAmp1 = numchop(AmpTetraInPrism(a,b,c,j1,j2,j3))
		TempAmp2 = numchop(AmpTetraInPrism(a,b,c,j4,j5,j6))
		if TempAmp1!=0 && TempAmp2!=0
			push!(Amp,TempAmp1*TempAmp2/(visqrt(a)*visqrt(b)*visqrt(c)))
			push!(Index,SuperIndex([a,b,c,j1,j2,j3,j4,j5,j6]))
		end
	end
	return Amp,Index
end

end # module


using Test, YaoPlots
using ZXCalculus, Multigraphs

@testset "zx plot" begin
    g = Multigraph(6)
    for e in [[1,3],[2,3],[3,4],[4,5],[4,6]]
        add_edge!(g, e)
    end
    ps = [0, 0, 0//1, 2//1, 0, 0]
    v_t = [SpiderType.In, SpiderType.Out, SpiderType.X, SpiderType.Z, SpiderType.Out, SpiderType.In]
    zxd = ZXDiagram(g, v_t, ps)
    plot(zxd)
    replace!(Rule{:b}(), zxd)
    plt = plot(zxd)
    @test plt !== nothing
end