# Author: Lukasz Bratos

# Wektory do testowania, odpowiednio dla Float64 i Float32 oraz prawidlowa odpowiedz
x64 = [2.718281828,-3.141592654,1.414213562,0.5772156649,0.3010299957]
y64 = [1486.2497,878366.9879,-22.37492,4773714.647,0.000185049]
x32 = map(a -> convert(Float32, a), x64)
y32 = map(a -> convert(Float32, a), y64)
answer = -1.00657107000000e-11

# Algorytm 'w przod'. Przyjmuje wektor x oraz y i rozmiar wektorow n.
function forward(x, y, n)
    s = 0.0
    for i in 1:n
        s = s + x[i] * y[i]
    end
    return s
end

# Algorytm 'w tyl'. Przyjmuje wektor x oraz y i rozmiar wektorow n.
function backward(x, y, n)
    s = 0.0
    for i in reverse(1:n)
        s = s + x[i] * y[i]
    end
    return s
end

# Algorytm 'od najwiekszego do najmniejszego'. Przyjmuje wektor x oraz y, typ zmiennopozycyjny w ktorym wykonywane sa obliczenia oraz rozmiar wektorow n.
function fromTheBiggestToTheSmallest(x, y, type, n)
    partials = Array{type}(undef, 5)
    for i in 1:n
        partials[i] = x[i] * y[i]
    end

    pos = sort(filter(x -> x > 0, partials), rev=true)
    neg = sort(filter(x -> x < 0, partials))

    return foldl(+, pos) + foldl(+, neg)
end

# Algorytm 'od najmniejszego do najwiekszego'. Przyjmuje wektor x oraz y, typ zmiennopozycyjny w ktorym wykonywane sa obliczenia oraz rozmiar wektorow n.
function fromTheSmallestToTheBiggest(x, y, type, n)
    partials = Array{type}(undef, 5)
    for i in 1:n
        partials[i] = x[i] * y[i]
    end

    pos = sort(filter(x -> x > 0, partials))
    neg = sort(filter(x -> x < 0, partials), rev=true)

    return foldl(+, pos) + foldl(+, neg)
end

println("Wynik dla 'w przod' i Float64: ", forward(x64, y64, 5))
println("Wynik dla 'w przod' i Float32: ", forward(x32, y32, 5))

println("Wynik dla 'w tyl' i Float64: ", backward(x64, y64, 5))
println("Wynik dla 'w tyl' i Float32: ", backward(x32, y32, 5))

println("Wynik dla najwiekszy-najmniejszy i Float64: ", fromTheBiggestToTheSmallest(x64, y64, Float64, 5))
println("Wynik dla najwiekszy-najmniejszy i Float32: ", fromTheBiggestToTheSmallest(x32, y32, Float32, 5))

println("Wynik dla najmniejszy-najwiekszy i Float64: ", fromTheSmallestToTheBiggest(x64, y64, Float64, 5))
println("Wynik dla najmniejszy-najwiekszy i Float32: ", fromTheSmallestToTheBiggest(x32, y32, Float32, 5))


module SmallLeastSquaresDemo

using VarianceReducedSG
using LinearAlgebra
using Random
using Test

# Square-norm gradient
sqgrad(x) = x

function setup_ls(N, n, mt)
	# Generate data
	A = randn(mt, n,N)
	x_star = randn(mt, N)

	y = A*x_star # No noise

	# Normalize and reformat
	norm_consts = [norm(A[i,:]) for i = 1:size(A,1)]
	data = [A[i,:]./norm_consts[i] for i = 1:size(A,1)]
	b = [y[i]/norm_consts[i] for i = 1:length(y)]

	# Initial guess
	x0 = randn(mt, N)

	return data, b, x0, x_star
end

function solve_and_test(alg, niter, loggers, nlogs, x_star)
	iterate!(alg, niter, loggers, nlogs)
	@test isapprox(primiterates(alg), x_star)

	iterate!(alg, niter)
	@test isapprox(primiterates(alg), x_star)

	iterate!(alg, niter, loggers, nlogs)
	@test isapprox(primiterates(alg), x_star)
end

################################################################################
# SAGA, SAG, SVAG
################################################################################
function SAGlike_test()
	N = 100
	n = 1000
	mt = MersenneTwister(123)

	data, b, x0, x_star = setup_ls(N,n,mt)

	# Gradients of square cost
	f(x,i) = sqgrad(x - b[i])
	f!(y,x,i) = (y .= data[i].*sqgrad(dot(data[i],x) - b[i]))

	# VR gradient
	vrg = VRGradient(f!, zeros(N), n)
	vrg_lin = LinearVRG(f, data, zeros(N))

	# Create Algorithms
	algs = [
		(SAGA(vrg, x0, 1/3, mt), "SAGA - VRGradient"),
		(SAGA(vrg_lin, x0, 1/3, mt), "SAGA - LinearVRG"),
		(SAGA(vrg_lin, x0, 1/3, mt, weights=.1 .+rand(mt,n)), "SAGA - LinearVRG - Importance"),
		(SAG(vrg, x0, 1/3, mt), "SAG - VRGradient"),
		(SAG(vrg_lin, x0, 1/3, mt), "SAG - LinearVRG"),
		(SAG(vrg, x0, 1/3, mt, weights=.1 .+rand(mt,n)), "SAG - VRGradient - Importance"),
		(SVAG(vrg, x0, 1/3, n/10, mt), "SVAG - VRGradient"),
		(SVAG(vrg_lin, x0, 1/3, n/10, mt), "SVAG - LinearVRG"),
		(SVAG(vrg_lin, x0, 1/3, n/10, mt, weights=.1 .+rand(mt,n)), "SVAG - LinearVRG - Importance"),
		(ASVAG(vrg, x0, 1/3, mt), "ASVAG - VRGradient"),
		(ASVAG(vrg_lin, x0, 1/3, mt), "ASVAG - LinearVRG"),
		(ASVAG(vrg_lin, x0, 1/3, mt, weights=.1 .+rand(mt,n)), "ASVAG - LinearVRG - Importance"),
		(ASVAG(vrg_lin, x0, 1/3, mt, enforce_step=true), "ASVAG - LinearVRG - Theoretical stepsize enforced"),
		]

	for (alg, descr) in algs
		# Create Loggers
		fprog(x,y) = norm(x - x_star)
		loggers = (
			ShowTime(),
			ShowFuncVal(fprog, "|| x - x^*||"),
			ShowNewLine("...")
			)

		println(descr)

		# Solve
		niter = n*500
		nlogs = 10
		solve_and_test(alg, niter, loggers, nlogs, x_star)
	end
end

################################################################################
# SVRG
################################################################################
function SVRG_test()
	N = 100
	n = 1000
	mt = MersenneTwister(123)

	data, b, x0, x_star = setup_ls(N,n,mt)

	# Gradients of square cost
	f(x,i) = sqgrad(x - b[i])
	f!(y,x,i) = (y .= data[i].*sqgrad(dot(data[i],x) - b[i]))

	# VR gradient
	vrg = VRGradient(f!, zeros(N), n)
	vrg_lin = LinearVRG(f, data, zeros(N))
	vrg_uni = UniformVRG(f!, zeros(N), n)

	# Create Algorithm
	stagelen = Int(ceil(n/2))
	algs = [
		(SVRG(vrg, x0, 1/3, stagelen, mt), "SVRG - VRGradient"),
		(SVRG(vrg_uni, x0, 1/3, stagelen, mt), "SVRG - UniformVRG"),
		(SVRG(vrg_lin, x0, 1/3, stagelen, mt), "SVRG - LinearVRG"),
		(SVRG(vrg_uni, x0, 1/3, stagelen, mt, weights=.1 .+rand(mt,n)), "SVRG - UniformVRG - Importance"),
		]

	for (alg, descr) in algs
		# Create Loggers
		fprog(x,y) = norm(x - x_star)
		loggers = (
			ShowFuncVal(fprog, "|| x - x^*||"),
			ShowIterations(),
			ShowNewLine()
			)

		println(descr)

		# Solve
		niter = 100*stagelen
		nlogs = 10
		solve_and_test(alg, niter, loggers, nlogs, x_star)
	end
end

################################################################################
# Loopless SVRG
################################################################################
function LSVRG_test()
	N = 100
	n = 1000
	mt = MersenneTwister(123)

	data, b, x0, x_star = setup_ls(N,n,mt)

	# Gradients of square cost
	f(x,i) = sqgrad(x - b[i])
	f!(y,x,i) = (y .= data[i].*sqgrad(dot(data[i],x) - b[i]))

	# VR gradient
	vrg = VRGradient(f!, zeros(N), n)
	vrg_lin = LinearVRG(f, data, zeros(N))
	vrg_uni = UniformVRG(f!, zeros(N), n)

	q = 10/n # Expected update frequency
	algs = [
		(LSVRG(vrg, x0, 1/3, q, mt), "LSVRG - VRGradient"),
		(LSVRG(vrg_uni, x0, 1/3, q, mt), "LSVRG - UniformVRG"),
		(LSVRG(vrg_lin, x0, 1/3, q, mt), "LSVRG - LinearVRG"),
		(LSVRG(vrg, x0, 1/3, q, mt, weights=.1 .+ rand(mt,n)), "LSVRG - VRGradient - Importance"),
		]

	for (alg, descr) in algs
		# Create Loggers
		fprog(x,y) = norm(x - x_star)
		loggers = (
			ShowIterations(),
			ShowFuncVal(fprog, "|| x - x^*||"),
			ShowNewLine()
			)

		println(descr)

		# Setup iterations counts
		niter = 1000/q
		nlogs = 10

		# Create Algorithm and solve
		solve_and_test(alg, niter, loggers, nlogs, x_star)
	end

	# Solve again for smaller expected update frequency
	q = .1/n
	niter = 100/q
	nlogs = 10

	# Create Loggers
	fprog(x,y) = norm(x - x_star)
	loggers = (
		ShowIterations(),
		ShowFuncVal(fprog, "|| x - x^*||"),
		ShowNewLine()
		)
	println("LSVRG - Low update frequency")
	alg = LSVRG(vrg, x0, 1/3, q, mt)
	solve_and_test(alg, niter, loggers, nlogs, x_star)
end

################################################################################
# Incoherent Loopless SVRG
################################################################################
function ILSVRG_test()
	N = 100
	n = 1000
	mt = MersenneTwister(123)

	data, b, x0, x_star = setup_ls(N,n,mt)

	# Gradients of square cost
	f(x,i) = sqgrad(x - b[i])
	f!(y,x,i) = (y .= data[i].*sqgrad(dot(data[i],x) - b[i]))

	# VR gradient
	vrg = VRGradient(f!, zeros(N), n)
	vrg_lin = LinearVRG(f, data, zeros(N))
	vrg_uni = UniformVRG(f!, zeros(N), n)

	q = 10/n # Expected update frequency
	algs = [
		(ILSVRG(vrg, x0, 1/3, q, mt), "ILSVRG - VRGradient"),
		(ILSVRG(vrg_lin, x0, 1/3, q/10, mt), "ILSVRG - LinearVRG"),
		(ILSVRG(vrg, x0, 1/3, 2*q, mt, weights=.1 .+ rand(mt,n)), "ILSVRG - VRGradient - Importance"),
		]

	for (alg, descr) in algs
		# Create Loggers
		fprog(x,y) = norm(x - x_star)
		loggers = (
			ShowIterations(),
			ShowFuncVal(fprog, "|| x - x^*||"),
			ShowNewLine()
			)

		println(descr)

		# Setup iterations counts
		niter = 1000/q
		nlogs = 10

		# Create Algorithm and solve
		solve_and_test(alg, niter, loggers, nlogs, x_star)
	end
end

################################################################################
# qSAGA
################################################################################
function QSAGA_test()
	N = 100
	n = 1000
	mt = MersenneTwister(123)

	data, b, x0, x_star = setup_ls(N,n,mt)

	# Gradients of square cost
	f(x,i) = sqgrad(x - b[i])
	f!(y,x,i) = (y .= data[i].*sqgrad(dot(data[i],x) - b[i]))

	# VR gradient
	vrg = VRGradient(f!, zeros(N), n)
	vrg_lin = LinearVRG(f, data, zeros(N))

	algs = [
		(QSAGA(vrg, x0, 1/3, 1, mt), "QSAGA - VRGradient - q=1"),
		(QSAGA(vrg_lin, x0, 1/3, 10, mt), "QSAGA - LinearVRG, q=10"),
		(QSAGA(vrg, x0, 1/3, 10, mt, weights=.1 .+ rand(mt,n)), "LSVRG - VRGradient - Importance - q=10"),
		]

	for (alg, descr) in algs
		# Create Loggers
		fprog(x,y) = norm(x - x_star)
		loggers = (
			ShowIterations(),
			ShowFuncVal(fprog, "|| x - x^*||"),
			ShowNewLine()
			)

		println(descr)

		# Setup iterations counts
		niter = 100*n
		nlogs = 10

		# Create Algorithm and solve
		solve_and_test(alg, niter, loggers, nlogs, x_star)
	end

	# Solve again but without replacement in the dual update
	niter = 50*n
	nlogs = 10

	# Create Loggers
	fprog(x,y) = norm(x - x_star)
	loggers = (
		ShowIterations(),
		ShowFuncVal(fprog, "|| x - x^*||"),
		ShowNewLine()
		)
	println("LSVRG - q = 20 - Without replacement")
	alg = QSAGA(vrg, x0, 1/3, 20, mt, replace=false)
	solve_and_test(alg, niter, loggers, nlogs, x_star)
	println("LSVRG - q = 20 - With replacement")
	alg = QSAGA(vrg, x0, 1/3, 20, mt, replace=true)
	solve_and_test(alg, niter, loggers, nlogs, x_star)
end

################################################################################
# Testsets
################################################################################
@testset "Least Squares" begin
	SAGlike_test()
	SVRG_test()
	LSVRG_test()
	ILSVRG_test()
	QSAGA_test()
end

end


"""
    Triplets(D::EDM{T}) where T

Compute the triplets from D.
"""
function TripletEmbeddings.Triplets(D::EDM{T}) where T <: Real

    n = size(D,1)
    triplets = Vector{Tuple{Int,Int,Int}}(undef, n*binomial(n-1, 2))
    counter = 0

    for k = 1:n, j = 1:k-1, i = 1:n
        if i != j && i != k
            if D[i,j] < D[i,k]
                counter +=1
                @inbounds triplets[counter] = (i, j, k)
            elseif D[i,j] > D[i,k]
                counter += 1
                @inbounds triplets[counter] = (i, k, j)
            end
        end
    end

    return Triplets(triplets[1:counter])
end

"""
    Triplets(D::MaskedEDM{T}) where T

Compute the triplets from D, using only its unmasked (or non-zero) values.
"""
function TripletEmbeddings.Triplets(D::MaskedEDM{T}) where T <: Real

    n = size(D,1)
    Dm = masked(D)
    triplets = Vector{Tuple{Int,Int,Int}}(undef, n*binomial(n-1, 2))
    counter = 0

    for k = 1:n, j = 1:k-1, i = 1:n
        if i != j && i != k
            if Dm[i,j] != 0 && Dm[i,k] != 0
                if Dm[i,j] < Dm[i,k]
                    counter +=1
                    @inbounds triplets[counter] = (i, j, k)
                elseif Dm[i,j] > Dm[i,k]
                    counter += 1
                    @inbounds triplets[counter] = (i, k, j)
                end
            end
        end
    end

    return Triplets(triplets[1:counter])
end


# ---
# title: 992. Subarrays with K Different Integers
# id: problem992
# author: Tian Jun
# date: 2020-10-31
# difficulty: Hard
# categories: Hash Table, Two Pointers, Sliding Window
# link: <https://leetcode.com/problems/subarrays-with-k-different-integers/description/>
# hidden: true
# ---
#
# Given an array `A` of positive integers, call a (contiguous, not necessarily
# distinct) subarray of `A` _good_ if the number of different integers in that
# subarray is exactly `K`.
#
# (For example, `[1,2,3,1,2]` has `3` different integers: `1`, `2`, and `3`.)
#
# Return the number of good subarrays of `A`.
#
#
#
# **Example 1:**
#
#
#
#     Input: A = [1,2,1,2,3], K = 2
#     Output: 7
#     Explanation: Subarrays formed with exactly 2 different integers: [1,2], [2,1], [1,2], [2,3], [1,2,1], [2,1,2], [1,2,1,2].
#
#
# **Example 2:**
#
#
#
#     Input: A = [1,2,1,3,4], K = 3
#     Output: 3
#     Explanation: Subarrays formed with exactly 3 different integers: [1,2,1,3], [2,1,3], [1,3,4].
#
#
#
#
# **Note:**
#
#   1. `1 <= A.length <= 20000`
#   2. `1 <= A[i] <= A.length`
#   3. `1 <= K <= A.length`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


using Doubles
using Base.Test

# courtesy: https://github.com/simonbyrne/DoubleDouble.jl

x = sqrt(2.0)
bx = big(x)
sx = x
dx = Double(x)

y = 0.1
by = big(y)
sy = y
dy = Double(y)

@test x == sx == dx
@test y == sy == dy

dxy = dx*dy
bxy = bx*by
@test sx*sy == dxy
@test x*y == Float64(dxy)
# @test dxy == Double(bxy)

@test x+y == Float64(dx+dy)
# @test dx+dy == Double(bx+by)

@test x-y == Float64(dx-dy)
# @test dx-dy == Double(bx-by)

@test x/y == Float64(dx/dy)
# @test dx/dy == Double(bx/by)

@test sqrt(y) == Float64(sqrt(dy))
# @test sqrt(dy) == Double(sqrt(by))


# https://www.hackerrank.com/challenges/30-conditional-statements/problem

n = parse(Int, readline())

function check_weird(v::Int)
    r_weird = "Weird"
    r_not_weird = "Not Weird"
    is_even = v % 2 == 0
    if is_even
        if 2 ≤ v ≤ 5
            return r_not_weird
        elseif 6 ≤ v ≤ 20
            return r_weird
        else
            return r_not_weird
        end
    else
        return r_weird
    end
end

println(check_weird(n))

Practice/Tutorials/30 Days of Code/Day 03 Intro to Conditional Statements.jl


function bubblesort!(x::AbstractArray)
    n = length(x)
    for i ∈ 1:n
        for j ∈ (i+1):n
            @inbounds if x[j] ≤ x[i]
                @inbounds tmp  = x[i]
                @inbounds x[i] = x[j]
                @inbounds x[j] = tmp
            end
        end
    end
    return nothing
end


## binning algorithms

struct EqualSize
    "Number of bins."
    n::Int

    @doc """
        EqualSize(; n::Int=10)

    Create binning algorithm with `n` bins of the same size.
    """
    function EqualSize(; n::Int=10)
        n > 0 || error("number of bins must be positive")
        return new(n)
    end
end

struct EqualMass
    "Number of bins."
    n::Int

    @doc """
        EqualMass(; n::Int=10)

    Create binning algorithm with `n` bins of (approximately) equal mass.
    """
    function EqualMass(; n::Int=10)
        n > 0 || error("number of bins must be positive")
        return new(n)
    end
end

function bins(::Any, ::Any, algorithm::EqualSize)
    return StatsBase.Histogram(range(0, 1; length=algorithm.n + 1))
end
function bins(probabilities, ::Any, algorithm::EqualMass)
    return StatsBase.Histogram(StatsBase.nquantile(probabilities, algorithm.n))
end

function binindex(bins::StatsBase.Histogram{<:Real,1}, x::Real)
    index = StatsBase.binindex(bins, x)
    nbins = length(bins.weights)
    edges = first(bins.edges)
    return if index == 0 && x == first(edges)
        1
    elseif index == nbins + 1 && x == last(edges)
        nbins
    else
        index
    end
end

function max_nonempty_bins(bins::StatsBase.Histogram{<:Real,1}, n::Int)
    return min(length(bins.weights), n)
end


# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.11.2
#   kernelspec:
#     display_name: Julia 1.8.0-DEV
#     language: julia
#     name: julia-1.8
# ---

# %%
using Distributions
using Plots

Distributions.Normal{T}(; μ, σ) where T = Normal{T}(μ, σ)
normals = [
    Normal{Float64}(μ=-0.42117172154916505, σ=1.2253594150160463)
    Normal{Float64}(μ=1.4244177594878171, σ=1.2253594150160463)
    Normal{Float64}(μ=-0.7646882557011875, σ=1.2253594150160463)
    Normal{Float64}(μ=-1.4994608515495296, σ=1.2253594150160463)
    Normal{Float64}(μ=0.1588834413196244, σ=1.2253594150160463)
    Normal{Float64}(μ=0.6906462431515679, σ=1.2253594150160463)
    Normal{Float64}(μ=-0.8610106914769075, σ=1.2253594150160463)
    Normal{Float64}(μ=0.39571181451559956, σ=1.2253594150160463)
    Normal{Float64}(μ=-0.7506648963208741, σ=1.2253594150160463)
    Normal{Float64}(μ=1.0347673088571538, σ=1.2253594150160463)
] |> x -> sort(x; by=mean)

ys = range(-10, 10; length=300)
m, n = length(normals), length(ys)
plot(legend=false)
for (x, normal) in zip(range(-1.5, 1.5; length=m), normals)
    plot3d!(fill(x, n), ys, pdf.(normal, ys))
end
plot!()

# %%
using LinearAlgebra
using Distributions
using Plots

dist_true = MvNormal([50, 50], 25^2*[1 1; 1 1.3])
n = 30
sample = rand(dist_true, n)
x, y = sample[1,:], sample[2,:]
X = x .^ (0:1)'
b̂ = X\y # b\hat TAB -> b̂
ŷ = X*b̂
u = norm(y - ŷ)/√(n - 2)
f(x) = evalpoly(x, b̂)
s(x) = u*√(1 + dot([1, f(x)], X'X\[1, f(x)]))
#disty(x) = Normal(f(x), s(x))
disty(x) = LocationScale(f(x), s(x), TDist(n-2))

plot(legend=false, xtick=-100:50:200, ytick=-100:50:200)
for xs in 0:10:100
    ys = range(-50, 150; length=200)
    plot3d!(fill(xs, length(ys)), ys, pdf.(disty(xs), ys))
end
scatter3d!(x, y, zeros(n); ms=2, msc=:auto, color=:red)
xs = range(-30, 130; length=100)
plot3d!(xs, f.(xs), zero(xs); color=:blue)

# %%


using Quantum, Test

@testset "Quantum" begin
    @testset "QuantumRegister" begin
        @testset "Construction" begin
            @test QuantumRegister(1) == [1, 0]
            @test QuantumRegister(2) == [1, 0, 0, 0]
            @test QuantumRegister(3) == [1, 0, 0, 0, 0, 0, 0, 0]
            @test QuantumRegister(2) == QuantumRegister(2)
            @test [1, 0, 0, 0] == QuantumRegister(2)
        end
        @testset "Length" begin
            @test length(QuantumRegister(1)) == 1
            @test length(QuantumRegister(2)) == 2
            @test length(QuantumRegister(3)) == 3
            @test length(QuantumRegister(4)) == 4
        end
    end
    @testset "Basic Gates" begin
        @testset "X" begin
            @test X(QuantumRegister(1), 1) == [0, 1]
            @test X(QuantumRegister(2), 1) == [0, 0, 1, 0]
            @test X(QuantumRegister(2), 2) == [0, 1, 0, 0]
            @test X(QuantumRegister(3), 2) == [0, 0, 1, 0, 0, 0, 0, 0]
        end
        @testset "H" begin
            @test H(QuantumRegister(1), 1) == [1 / √2, 1 / √2]
            @test H(X(QuantumRegister(1), 1), 1) == [1 / √2, -1 / √2]
        end
        @testset "CNOT" begin
            @test CNOT(QuantumRegister(2), 1, 2) == [1, 0, 0, 0]
            @test CNOT(X(QuantumRegister(2), 1), 1, 2) == [0, 0, 0, 1]
            @test CNOT(X(QuantumRegister(2), 2), 1, 2) == [0, 1, 0, 0]
            @test CNOT(X(QuantumRegister(2), 1), 2, 1) == [0, 0, 1, 0]
            @test CNOT(X(QuantumRegister(2), 2), 2, 1) == [0, 0, 0, 1]
            @test CNOT(X(QuantumRegister(3), 1), 1, 3) == [0, 0, 0, 0, 0, 1, 0, 0]
        end
        @testset "Swap" begin
            @test Swap(QuantumRegister(2), 1, 2) == [1, 0, 0, 0]
            @test Swap(QuantumRegister(2), 2, 1) == [1, 0, 0, 0]
            @test Swap(X(QuantumRegister(2), 1), 1, 2) == [0, 1, 0, 0]
            @test Swap(X(QuantumRegister(2), 2), 1, 2) == [0, 0, 1, 0]
            @test Swap(X(QuantumRegister(2), 1), 2, 1) == [0, 1, 0, 0]
        end
        @testset "Entanglement" begin
            @test CNOT(H(QuantumRegister(2), 1), 1, 2) == [1 / √2, 0, 0, 1 / √2]
            @test CNOT(H(X(QuantumRegister(2), 2), 1), 1, 2) == [0, 1 / √2, 1 / √2, 0]
        end
    end
    @testset "Single Qubit" begin
        @testset "Pauli Gates" begin
            @test Y(QuantumRegister(1), 1) == [0.0, im]
            @test Y(X(QuantumRegister(1), 1), 1) == [-im, 0.0]
            @test Z(QuantumRegister(1), 1) == [1, 0]
            @test Z(X(QuantumRegister(1), 1), 1) == [0, -1]
        end
        @testset "S(dag)" begin
            @test S(QuantumRegister(1), 1) == [1, 0]
            @test S(X(QuantumRegister(1), 1), 1) == [0, im]
            @test S(QuantumRegister(1), 1) == T(T(QuantumRegister(1), 1), 1)
            @test Sdag(QuantumRegister(1), 1) == [1, 0]
            @test Sdag(X(QuantumRegister(1), 1), 1) == [0, -im]
        end
        @testset "T(dag)" begin
            @test T(QuantumRegister(1), 1) == [1, 0]
            @test T(X(QuantumRegister(1), 1), 1) == [0, exp(im * pi / 4)]
            @test T(T(QuantumRegister(1), 1), 1) == S(QuantumRegister(1), 1)
            @test Tdag(QuantumRegister(1), 1) == [1, 0]
            @test Tdag(X(QuantumRegister(1), 1), 1) == [0, exp(-im * pi / 4)]
        end
    end
    @testset "Controlled Gates" begin
        @testset "Single Control" begin
            @test CZ(QuantumRegister(2), 1, 2) == [1, 0, 0, 0]
            @test CZ(X(QuantumRegister(2), 1), 1, 2) == [0, 0, 1, 0]
            @test CZ(X(QuantumRegister(2), 2), 1, 2) == [0, 1, 0, 0]
            @test CZ(X(QuantumRegister(2), 1), 2, 1) == [0, 0, 1, 0]
            @test CZ(X(QuantumRegister(2), 2), 2, 1) == [0, 1, 0, 0]
            @test CZ(X(QuantumRegister(3), 1), 1, 3) == [0, 0, 0, 0, 1, 0, 0, 0]
            @test CNOT(H(QuantumRegister(2), 2), 2, 1) == [1 / √2, 0, 0, 1 / √2]
            @test CNOT(H(X(QuantumRegister(2), 1), 2), 2, 1) == [0, 1 / √2, 1 / √2, 0]
        end
    end
end


#reload("MetabolicFluxAnalysis.jl")
using MetabolicFluxAnalysis
using CoordinateTransformations

# current tests..
fluxi_1_f = open("testdata/test_01.fluxi")
fluxi_1_s = readstring(fluxi_1_f)
fluxi1 = MetabolicFluxAnalysis.parseFluxi( fluxi_1_s )

# now get and then parametrize the flux polytope
fp = MetabolicFluxAnalysis.get_flux_polyhedron( fluxi1.net , fluxi1.linear_constraints["ConstraintsExchange"] , max_val=200 )

# compute transform to make full-dimensional:
( fp_fd , at_fd ) = MetabolicFluxAnalysis.find_minimal_affine_subspace(fp)
fp_fd_a = MetabolicFluxAnalysis.remove_redundant_constraints( fp_fd )
( fp_rounded , at_rounded ) = MetabolicFluxAnalysis.round_polytope(fp_fd_a)

at_opt2flux = CoordinateTransformations.compose(at_fd , at_rounded)

# test chebyshev comp.
x_int_cheby = MetabolicFluxAnalysis.solve_chebyshev(fp_fd.G,fp_fd.h[:])

@test all( fp_fd.G * x_int_cheby .< fp_fd.h )
print( fp_fd.h - fp_fd.G * x_int_cheby  )

## test with toy polytopes..
d_test = 100
g_test = [eye(d_test);-eye(d_test)]; h_test = 10*rand(2*d_test,1);
a_test = rand(2,d_test)       ; b_test = 0.1 * rand(2,1);
f_test = MetabolicFluxAnalysis.Polyhedron(g_test,h_test,a_test,b_test)
rfd_parametrization = MetabolicFluxAnalysis.get_rounded_fulldimensional_parametrization( f_test ; remove_red_constr=1)
@assert MetabolicFluxAnalysis.isinside( rfd_parametrization.p_opt , rfd_parametrization.cheby_p_opt[:,:] )

@test MetabolicFluxAnalysis.isinside( rfd_parametrization.p_opt , rfd_parametrization.cheby_p_opt[:,:] )

# generate some unifrom samples over polytope:
xx = MetabolicFluxAnalysis.sample_unif_hr( rfd_parametrization.p_opt , rfd_parametrization.cheby_p_opt , 1000 ,1000 );

xx_raw = MetabolicFluxAnalysis.eval_opt2raw(rfd_parametrization,xx)

@test MetabolicFluxAnalysis.checksamples( rfd_parametrization.p_opt , xx     ; eq_tol=1e-4)
@test MetabolicFluxAnalysis.checksamples( rfd_parametrization.p_raw , xx_raw ; eq_tol=1e-4)

##

# some more basic tests
# check that things behave correctly for unbounded polyhedra
using MetabolicFluxAnalysis
bb_01 = MetabolicFluxAnalysis.bounding_box([-10.;-10;-10],[10.;10;20])
unbounded_polyhedron = MetabolicFluxAnalysis.Polyhedron( bb_01.G[1:1,:],bb_01.h[1:1,:],bb_01.A[:,:],bb_01.b[:,:] )
lpm_result = MetabolicFluxAnalysis.solve_lp_max_v(unbounded_polyhedron,[0.;1;0])

#(mve_f,mve_E) =  MetabolicFluxAnalysis.mve_run(fp_fd_a.G,fp_fd_a.h[:])

# test rounding transform:
cube_g = [1.0*eye(4);-eye(4)]
cube_h = 1.0*collect(1:8)

p_c = MetabolicFluxAnalysis.Polyhedron(cube_g,cube_h[:,:])

# remove_redundant_constraintss
param1= fp_fd_a.G
param2 = fp_fd_a.h[:]
param3 = [79.514 ;52.79]
(mve_f,mve_E) =  MetabolicFluxAnalysis.mve_solver( param1,param2,param3 )

(mve_f,mve_E) =  MetabolicFluxAnalysis.mve_run(cube_g,cube_h)

# transform somehow:

p_fixed = MetabolicFluxAnalysis.Polyhedron(fp_fd,MetabolicFluxAnalysis.bounding_box([-1000.;-1000],[1000.;1000]) )

(mve_f,mve_E) = MetabolicFluxAnalysis.mve_run( p_fixed.G , p_fixed.h[:] )
(mve_f,mve_E) = MetabolicFluxAnalysis.mve_run( fp_fd.G , fp_fd.h[:] )

# compute rounding transform:
(mve_f,mve_E) = MetabolicFluxAnalysis.mve_run( fp_fd.G , fp_fd.h[:] )


module CP
include("dynam_anim.jl")

using Distributions,Luxor

struct Params
  λ::Float64
end

getdefault(xs,i,def) = checkbounds(Bool,xs,i) ? @inbounds(xs[i]) : def

rate(σ::Vector{Bool},p::Params,i::Int)::Float64 = σ[i] ? 1 : p.λ * sum(getdefault(σ,i+j,0) for j in [-1,1])

rates(σ::Vector{Bool},p::Params) = [rate(σ,p,i) for i in eachindex(σ)]

function step(σ::Vector{Bool},p::Params)::Float64
  rs = rates(σ,p)
  r = sum(rs)
  if r == 0
    return Inf
  end
  p = rs/r
  i = rand(Categorical(p))
  σ[i] = !σ[i]
  rand(Exponential(1/r))
end

function setup(σ::Vector{Bool},w,h)
  Tiler(w,h,1,length(σ),margin=0)
end

function draw(σ::Vector{Bool},tiles)
  for (pos,n) in tiles
    setcolor(σ[n] ? "red" : "green")
    box(pos, tiles.tilewidth, tiles.tileheight, :fill)
  end
end

function contact_anim(init,p,frames,steps_per_frame,width,height,path)
  ds = DynamAnim.DynamicalSystem(init,p,step)
  ap = DynamAnim.AnimParams(setup,draw,frames,steps_per_frame,width,height,"Contact Process")
  DynamAnim.anim(ds,ap,path)
end

end


export ActivateFuncs

module ActivateFuncs
import Base.sign
using QuaternionArgs
using Quaternions

export sign, quantization

# split type sign functions {{{
function sign{T<:Real}(z::Complex{T})
    return complex(sign(real(z)) ,sign(imag(z)))
end

function sign{T<:Real}(q::Quaternion{T})
    return Quaternion{T}(sign(q.q0),sign(q.q1),sign(q.q2),sign(q.q3))
end
# }}}

# quantization functions for multistate neurons {{{
function quantization{T<:Real}(z::Complex{T},K::Integer)
    halfphi = pi/K
    phi = 2.0*halfphi
    theta = arg(z) + halfphi
    q = floor(theta / phi)
    w = q*phi
    return complex(cos(w),sin(w))
end
function arg{T<:Real}(z::Complex{T})
    return atan2(imag(z),real(z))
end

function quantization_{T<:Real}(q::Quaternion{T},A::Integer,B::Integer,C::Integer)
    qarg = QuaternionArg(q)
    if randbool()
        phi = qarg.phi
        theta = qsign(qarg.theta, 0.5, B)
        psi = qsign(qarg.psi, 0.25, C)
    else
        phi = qsign(qarg.phi, 1.0, A)
        theta = qsign(qarg.theta, 0.5, B)
        psi = qarg.psi
    end
    return Quaternion(QuaternionArg(one(T),phi,theta,psi))
end

function quantization{T<:Real}(q::Quaternion{T},A::Integer,B::Integer,C::Integer)
    qarg = QuaternionArg(q)

    phi = qsign(qarg.phi, 1.0, A)
    theta = qsign(qarg.theta, 0.5, B)
    psi = qsign(qarg.psi, 0.25, C)

    return Quaternion(QuaternionArg(one(T),phi,theta,psi))
end

function qsign(phase::Real, coef::Real, K::Integer)
    dphase = 2*pi/K*coef
    phase0 = -pi*coef
    flg = false
    p = phase0

    # phase regulation
    phase = phase < phase0 ? phase0 : (phase >= -phase0 ? -phase0-0.5*dphase : phase)
    for i=1:K
        if p <= phase && phase < p+dphase
            flg = true
            break
        end
        p = phase0+dphase*i
    end
    @assert(flg)

    return p+0.5*dphase
end
# }}}

end


#=

  TSP (traveling salesperson problem) in Julia ConstraintSolver.jl

  Inspired by the code from lecture notes
  Ulf Nilsson: Transparencies for the course TDDD08 Logic
  Programming, page 6f
  http://www.ida.liu.se/~TDDD08/misc/ulfni.slides/oh10.pdf

  Model created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Julia page: http://www.hakank.org/julia/

=#
using ConstraintSolver, JuMP
using Cbc, GLPK, Ipopt
const CS = ConstraintSolver
include("constraints_utils.jl")

function tsp(problem,print_solutions=true,all_solutions=true)

    cbc_optimizer = optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0)
    glpk_optimizer = optimizer_with_attributes(GLPK.Optimizer)
    ipopt_optimizer = optimizer_with_attributes(Ipopt.Optimizer)

    model = Model(optimizer_with_attributes(CS.Optimizer,   "all_solutions"=> all_solutions,
                                                            # "all_optimal_solutions"=>true,
                                                            "logging"=>[],

                                                            "traverse_strategy"=>:BFS,
                                                            # "traverse_strategy"=>:DFS,
                                                            # "traverse_strategy"=>:DBFS,

                                                            "branch_split"=>:Smallest,
                                                            # "branch_split"=>:Biggest,
                                                            # "branch_split"=>:InHalf,

                                                            # "simplify"=>false,
                                                            # "simplify"=>true, # default

                                                            "time_limit"=>3,

                                                            # "lp_optimizer" => cbc_optimizer,
                                                            "lp_optimizer" => glpk_optimizer,
                                                            # "lp_optimizer" => ipopt_optimizer,
                                        ))

    problem = resize_matrix(problem)
    n, _ = size(problem)
    min_dist, max_dist = extrema(problem)
    println("min_dist:$min_dist max_dist:$max_dist (max_dist*n: $(max_dist*n))")

    @variable(model, 1 <= cities[1:n] <= n, Int)
    @variable(model, 1 <= path[1:n] <= n, Int)
    @variable(model, min_dist <= costs[1:n] <= max_dist, Int)
    @variable(model, 0 <= cost <= max_dist*n, Int)

    @constraint(model, cities in CS.AllDifferentSet())
    @constraint(model, path in CS.AllDifferentSet())
    circuit_path(model,cities,path)
    @constraint(model, cities[n] == 1)

    for i in 1:n
        my_element(model,cities[i],problem[i,:],costs[i])
    end
    @constraint(model, cost == sum(costs))

    # This is too slow!
    # Instead we check all solutions "manually" below and pick the best.
    # @objective(model,Min, cost)

    # Solve the problem
    println("solve")
    optimize!(model)

    status = JuMP.termination_status(model)
    # println("status:$status")
    if status == MOI.OPTIMAL
        num_sols = MOI.get(model, MOI.ResultCount())
        println("num_sols:$num_sols\n")
        min_cost = 99999999
        min_cities = []
        min_path = []
        min_costs = []
        if print_solutions
            for sol in 1:num_sols
                # println("solution #$sol")
                citiesx = convert.(Integer,JuMP.value.(cities; result=sol))
                costsx = convert.(Integer,JuMP.value.(costs; result=sol))
                costx = convert.(Integer,JuMP.value.(cost; result=sol))
                pathx = convert.(Integer,JuMP.value.(path; result=sol))
                if costx < min_cost
                    println("solution #$sol")
                    println("cities:$citiesx")
                    println("path  :$pathx")
                    println("costs :$costsx")
                    println("cost  :$costx")
                    println()
                    min_cost = costx
                    min_cities = citiesx
                    min_costs = costsx
                end

            end
        end
    else
        println("status:$status")
    end

    return status
end

tsp_problems = Dict(

#
# Problem from Nilsson cited above.
#
:nilsson =>
       [[ 0, 4, 8,10, 7,14,15],
        [ 4, 0, 7, 7,10,12, 5],
        [ 8, 7, 0, 4, 6, 8,10],
        [10, 7, 4, 0, 2, 5, 8],
        [ 7,10, 6, 2, 0, 6, 7],
        [14,12, 8, 5, 6, 0, 5],
        [15, 5,10, 8, 7, 5, 0]],

#
# This problem is from the SICStus example
# ./library/clpfd/examples/tsp.pl
# The "chip" examples
#
:chip =>
       [[0,205,677,581,461,878,345],
        [205,0,882,427,390,1105,540],
        [677,882,0,619,316,201,470],
        [581,427,619,0,412,592,570],
        [461,390,316,412,0,517,190],
        [878,1105,201,592,517,0,691],
        [345,540,470,570,190,691,0]],

# This problem is from the SICStus example
# ./library/clpfd/examples/tsp.pl
# The "ilog" examples
#
:ilog =>
       [[2,4,4,1,9,2,4,4,1,9],
        [2,9,5,5,5,2,9,5,5,5],
        [1,5,2,3,3,1,5,2,3,3],
        [2,6,8,9,5,2,6,8,9,5],
        [3,7,1,6,4,3,7,1,6,4],
        [1,2,4,1,7,1,2,4,1,7],
        [3,5,2,7,6,3,5,2,7,6],
        [2,7,9,5,5,2,7,9,5,5],
        [3,9,7,3,4,3,9,7,3,4],
        [4,1,5,9,2,4,1,5,9,2]],

# This problem is from
# GLPK:s example tsp.mod
# (via http://www.hakank.org/minizinc/tsp.mzn)
# """
# These data correspond to the symmetric instance ulysses16 from:
# Reinelt, G.: TSPLIB - A travelling salesman problem library.
# ORSA-Journal of the Computing 3 (1991) 376-84;
# http://elib.zib.de/pub/Packages/mp-testdata/tsp/tsplib
#
# The optimal solution is 6859
# """
:glpk =>
  [[0,509,501,312,1019,736,656,60,1039,726,2314,479,448,479,619,150],
   [509,0,126,474,1526,1226,1133,532,1449,1122,2789,958,941,978,1127,542],
   [501,126,0,541,1516,1184,1084,536,1371,1045,2728,913,904,946,1115,499],
   [312,474,541,0,1157,980,919,271,1333,1029,2553,751,704,720,783,455],
   [1019,1526,1516,1157,0,478,583,996,858,855,1504,677,651,600,401,1033],
   [736,1226,1184,980,478,0,115,740,470,379,1581,271,289,261,308,687],
   [656,1133,1084,919,583,115,0,667,455,288,1661,177,216,207,343,592],
   [60,532,536,271,996,740,667,0,1066,759,2320,493,454,479,598,206],
   [1039,1449,1371,1333,858,470,455,1066,0,328,1387,591,650,656,776,933],
   [726,1122,1045,1029,855,379,288,759,328,0,1697,333,400,427,622,610],
   [2314,2789,2728,2553,1504,1581,1661,2320,1387,1697,0,1838,1868,1841,1789,2248],
   [479,958,913,751,677,271,177,493,591,333,1838,0,68,105,336,417],
   [448,941,904,704,651,289,216,454,650,400,1868,68,0,52,287,406],
   [479,978,946,720,600,261,207,479,656,427,1841,105,52,0,237,449],
   [619,1127,1115,783,401,308,343,598,776,622,1789,336,287,237,0,636],
   [150,542,499,455,1033,687,592,206,933,610,2248,417,406,449,636,0]
  ],

)

print_solutions = true
all_solutions = true
# #=
solved = []
for problem in sort(collect(keys(tsp_problems)))
    println("\nproblem $problem")
    @time status = tsp(tsp_problems[problem],print_solutions,all_solutions)
    if status == MOI.OPTIMAL
        push!(solved,problem)
    end
end
# =#
# @time tsp(tsp_problems[:nilsson],print_solutions,all_solutions)
println("solved: $solved")


include("MackeyGlass.jl")

X, T = MGGenerator();


using Random
using Flux

struct Param
    beta
    lat
    nd
    volume
    tau
    nstep
    dt
    ntraj
    nrun
    nprint
    seed
    randinit
end

function Param(;
        beta = 6.0, lat = [64, 64],
        tau = 2.0, nstep = 50, ntraj = 256, nrun = 4, nprint = 256,
        seed = 11*13, randinit = false)
    nd = length(lat)
    volume = prod(lat)
    dt = tau / nstep
    Param(beta,lat,nd,volume,tau,nstep,dt,ntraj,nrun,nprint,seed,randinit)
end

summary(param::Param) = """latsize = $(param.lat)
    volume = $(param.volume)
    beta = $(param.beta)
    trajs = $(param.ntraj)
    tau = $(param.tau)
    steps = $(param.nstep)
    seed = $(param.seed)
    """
function uniquestr(param::Param)
    lat = join([string(x) for x in param.lat], ".")
    "out_l$(lat)_b$(param.beta)_n$(param.ntraj)_t$(param.tau)_s$(param.nstep).out"
end

action(param, f) = (-param.beta)*sum(cos.(plaqphase(f)))
force(param, f) = gradient(x -> action(param, x), f)[1]

plaqphase(f) = f[1,:,:] - f[2,:,:] - circshift(f[1,:,:], [0,-1]) + circshift(f[2,:,:], [-1,0])

topocharge(f) = round(sum(regularize(plaqphase(f))) / 2pi)
function regularize(f)
    p2 = 2pi
    # f_ = f .- pi
    p2 .* ((f .- pi) ./ p2 .- fld.((f .- pi), p2) .- 0.5)
end

function leapfrog(param, x, p)
    dt = param.dt
    x_ = x + 0.5dt .* p
    p_ = p + (-dt) .* force(param, x_)
    for i = 1:param.nstep
        x_ += dt .* p_
        p_ += (-dt) .* force(param, x_)
    end
    x_ += 0.5dt .* p_
    (x_, p_)
end
function hmc(param, x)
    p = randn(eltype(x), size(x))
    act0 = action(param, x) + 0.5*sum(p .* p)
    x_, p = leapfrog(param, x, p)
    x_ = regularize(x_)
    act = action(param, x_) + 0.5*sum(p .* p)
    prob = rand()
    dH = act-act0
    exp_mdH = exp(-dH)
    acc = prob < exp_mdH
    newx = acc ? x_ : x
    (dH, exp_mdH, acc, newx)
end

put(s) = write(stdout, s)

function main()
    param = Param(
        beta = 6.0, # 4.0
        lat = [64, 64], # [8, 8]
        tau = 2, # 0.3
        nstep = 50, # 3
        ntraj = 256, # 2**16 # 2**10 # 2**15
        nprint = 256,
        seed = 1331)

    Random.seed!(param.seed)

    field = param.randinit ? randn(Float64, (param.nd, param.lat...)) : zeros(Float64, (param.nd, param.lat...))  # mu, x, t

    open(uniquestr(param), "w") do O
        params = summary(param)
        write(O, params)
        put(params)
        plaq_, topo_ = (action(param, field) / (-param.beta*param.volume), topocharge(field))
        status = "Initial configuration:  plaq: $(plaq_)  topo: $(topo_)\n"
        write(O, status)
        put(status)
        ts = Float64[]
        for n = 1:param.nrun
            t = -time()
            for i = 1:param.ntraj
                dH, exp_mdH, acc, field = hmc(param, field)
                plaq = action(param, field) / (-param.beta*param.volume)
                topo = topocharge(field)
                ifacc = acc ? "ACCEPT" : "REJECT"
                status = "Traj: $((n-1)*param.ntraj+i) $ifacc:  dH: $dH  exp(-dH): $exp_mdH  plaq: $plaq  topo: $topo\n"
                write(O, status)
                if i % div(param.ntraj, param.nprint) == 0
                    put(status)
                end
            end
            t += time()
            append!(ts, t)
        end
        print("Run times: $ts\n")
        print("Per trajectory: $([t/param.ntraj for t in ts])\n")
    end
end

main()


"""
`pce(yr,r,g)`: Regional shares of final consumption
```math
\\alpha_{yr,r,g}^{pce} = \\dfrac{\\bar{pce}_{yr,r,g}}{\\sum_{r'} \\bar{pce}_{yr,r',g}}
```
"""
function share_pce!(d::Dict)
    println("  Calculating pce(yr,r,g), regional shares of final consumption")
    d[:pce] /= transform_over(d[:pce], :r)
    verify_over(d[:pce], :r) !== true && @error("PCE shares don't sum to 1.")
    return dropnan!(d[:pce])
end


# Use F.2 to generate them
function hermite_based_equilibrium(q, f)
    ρ = sum(f)
    u = velocity(q, f)
    T = temperature(q, f, ρ, u)

    hermite_based_equilibrium(q, ρ, u, T)
end

function hermite_based_equilibrium!(q::Q, ρ, u, T, f) where {Q <: Quadrature}
    N = div(LatticeBoltzmann.order(q), 2)

    # NOTE: we skip H_0 since its one (also because of Julia's 1 based index....)
    #Hs = [[hermite(n, q.abscissae[:, i], q) for i = 1:length(q.weights)] for n = 1:N]

    cs = 1 / q.speed_of_sound_squared
    D = dimension(q)
    δ(α, β) = α == β ? 1 : 0

    a0 = ρ
    H1 = [hermite(Val{1}, q.abscissae[:, i], q) for i in 1:length(q.weights)]
    a1 = u

    Hs = [[hermite(Val{n}, q.abscissae[:, i], q) for i in 1:length(q.weights)] for n in 1:N]
    a_eq = [equilibrium_coefficient(Val{n}, q, ρ, u, T) for n in 1:N]

    @inbounds for f_idx in 1:length(f)
        f[f_idx] =
            q.weights[f_idx] *
            (ρ + sum([dot(a_eq[n], Hs[n][f_idx]) / (factorial(n) * cs^n) for n in 1:N]))
    end
    return
end

# equilibrium(Val{1}, ρ, u, T)

function equilibrium_coefficient(::Type{Val{0}}, q::Quadrature, ρ, u, T)
    return ρ
end

function equilibrium_coefficient(::Type{Val{1}}, q::Quadrature, ρ, u, T)
    return ρ * u
end

function equilibrium_coefficient(::Type{Val{2}}, q::Quadrature, ρ, u, T)
    cs = 1 / q.speed_of_sound_squared
    return ρ * (u * u' + cs * (T - 1) * I)
end
function equilibrium_coefficient(::Type{Val{3}}, q::Quadrature, ρ, u, T)
    cs = 1 / q.speed_of_sound_squared
    D = dimension(q)
    return ρ * [
        u[a] * u[b] * u[c] + cs * (T - 1) * (
            u[a] * δ(b, c) + u[b] * δ(a, c) + u[c] * δ(a, b)
        )
        for a in 1:D, b in 1:D, c in 1:D
    ]
end
function equilibrium_coefficient(::Type{Val{4}}, q::Quadrature, ρ, u, T)
    cs = 1 / q.speed_of_sound_squared
    D = dimension(q)
    return ρ * [
        u[a] * u[b] * u[c] * u[d] +
        cs *
        (T - 1) *
        (
            u[a] * u[b] * δ(c, d) +
            u[a] * u[c] * δ(b, d) +
            u[a] * u[d] * δ(b, d) +
            u[b] * u[c] * δ(a, d) +
            u[b] * u[d] * δ(a, d) +
            u[c] * u[d] * δ(a, b)
        ) +
        cs^2 * (T - 1)^2 * (δ(a, b) * δ(c, d) + δ(a, c) * δ(b, d) + δ(a, d) * δ(b, c))
        for a in 1:D, b in 1:D, c in 1:D, d in 1:D
    ]
end


using Statistics, FixedEffectModels, DataFrames

struct EqmData
    Y::Array{Float64}
    D::Array{Float64}
    S::Array{Float64}
    Z::Array{Float64}
    W::Array{Bool, 1}
    U::Array{Float64}
    p::Float64
end

function EqmData(Y, D, S, W, U; p=0.5)
    return EqmData(Y, D, S, D .- S, W, U, p)
end

struct Partials
    yp1::Float64
    yp0::Float64
    dp1::Float64
    dp0::Float64
    sp1::Float64
    sp0::Float64
    μzp::Float64
end

function Partials(Y, D, S, W, U; p=0.5)
    ∂yp1, ∂yp0 = deriv_estimate(Y, W, U)
    ∂dp1, ∂dp0 = deriv_estimate(D, W, U)
    ∂sp1, ∂sp0 = deriv_estimate(S, W, U)
    ∂μzp = p*(∂dp1 - ∂sp1) + (1 - p)*(∂dp0 - ∂sp0)
    return Partials(∂yp1, ∂yp0, ∂dp1, ∂dp0, ∂sp1, ∂sp0, ∂μzp)
end

struct DiffMeans
    y::Float64
    d::Float64
    s::Float64
    z::Float64
end

function DiffMeans(Y, D, S, W; p=0.5)
    y = horvitz_thompson(Y, W)
    d = horvitz_thompson(D, W)
    s = horvitz_thompson(S, W)
    DiffMeans(y, d, s, d - s)
end

function horvitz_thompson(Y, W; p=0.5)
    data = (W./p - (1 .-W)./(1 - p)).*Y
    return mean(data)
end

function deriv_estimate(Y, W, U)
    ∂y1 = cov(Y[W.==1], U[W.==1])/var(U[W.==1])
    ∂y0 = cov(Y[W.==0], U[W.==0])/var(U[W.==0])
    return (∂y1, ∂y0)
end

function varianceDirect(Y, Z, W, ∂::Partials; p=0.5)
    a = (∂.yp1 - ∂.yp0)
    H = ∂.μzp
    A = -1*p .* W.*a./H.*Z + ((1 - p)*a/H) .* (1 .- W).*Z
    data = (W./p - (1 .-W)./(1 - p)).*(Y.+A)
    return std(data)/sqrt(length(Y))
end

function estimate_τ(data::EqmData)
    ∂ = Partials(data.Y, data.D, data.S, data.W, data.U)
    Δ = DiffMeans(data.Y, data.D, data.S, data.W)
    τD = Δ.y
    #println((data.p*∂.yp1 + (1- data.p)*∂.yp0))
    τI = -1.0*(data.p*∂.yp1 + (1- data.p)*∂.yp0)*Δ.z/∂.μzp
    σD = varianceDirect(data.Y, data.Z, data.W, ∂)
    σI = varianceIndirect(data.Y, data.Z, data.W, data.U, ∂, Δ)
    return (τD = τD, σD = σD, τI = τI, σI = σI)
end

function ivCOV(Y, X, Z)
    df = DataFrame(Y=Y, X=X, Z=Z)
    regmodel = reg(df, @formula(Y ~ (X ~ Z)), Vcov.robust())
    return (coef(regmodel)[2], stderror(regmodel)[2], coef(regmodel)[1])
end
function varianceIndirect(Y, Z, W, U, ∂::Partials, Δ::DiffMeans; p = 0.5)
    dpstar = -1.0*Δ.z
    coefIV, covIV, int = ivCOV(Y, Z, U)
    a = (∂.yp1 - ∂.yp0)
    H = ∂.μzp
    A = -1*p .* W.*a./H.*Z + ((1 - p)*a/H) .* (1 .- W).*Z
    data = (W./p - (1 .-W)./(1 - p)).*(Z.+A)
    covZ = var(data)/length(data)
    #σ2 = dpstar*(covIV^2*dpstar) + coefIV^2*covZ
    #println("Cov direct:", sqrt(dpstar*(covIV^2*dpstar)))
    ehat = Y .- Z.* coefIV .- int
    estimatehat = mean(ehat.^2)*(Δ.z/∂.μzp)^2/(length(Y)*U[1]^2)
    σ2 = estimatehat + coefIV^2*covZ
    return sqrt(σ2)#/abs(U[1])
end


"""The minimum mass considered."""
const mlow = 3.0

"""The maximum mass considered."""
const mhigh = 20.0

"""The reference redshift, at which rates are evaluated."""
const zref = 0.4

function _fm(m, a1, a2, mb)
    x = m/mb
    if m < mlow || m > mhigh
        -Inf
    elseif m < mb
        a1*log(x)
    else
        a2*log(x)
    end
end

"""
    make_log_dN(a1, a2, mb, mu_q, sigma_q)

Produces a function that computes `(m1, m2) -> log_dNdm1dm2` for the population
model.

The population model is

```math
m_1 m_2 \\frac{\\mathrm{d}N}{\\mathrm{d} m_1 \\mathrm{d} m_2 \\mathrm{d} V \\mathrm{d} t} = R f_m\\left( m_1 \\right) f_m\\left( m_2 \\right) \\exp\\left( -\\frac{\\left(q - \\mu_q\\right)^2}{2 \\sigma_q^2} + \\frac{\\left(1-\\mu_q\\right)^2}{2 \\sigma_q^2} \\right)^2
```
where
```math
f_m\\left( m \\right) = \\begin{cases}
m^{\\alpha_1} & m_\\mathrm{min} < m < m_b \\
m^{\\alpha_2} & m_b \\leq m < m_\\mathrm{max}
\\end{cases}
```
is the "common" black hole mass function, and a Gaussian-shaped pairing function
in ``q`` is used with parameters ``mu_q`` and ``sigma_q``.  The function
returned is normalized so that ``R`` measures the volumetric merger rate per
natural log mass squared at `m1 = m2 = m_b`.
"""

square(x) = x*x
function make_log_dN(a1, a2, mb, mu_q, sigma_q)
    function log_dN(m1, m2)
        q = m2/m1
        _fm(m1, a1, a2, mb) + _fm(m2, a1, a2, mb) - log(m1) - log(m2) - 0.5*square((q-mu_q)/sigma_q) + 0.5*square((1-mu_q)/sigma_q)
    end
    log_dN
end

"""
    make_dNdm1(a1, a2, mb, mu_q, sigma_q)

Returns a function that computes the marginal density (unnormalized) in `m1`.
"""
function make_dNdm1(a1, a2, mb, mu_q, sigma_q)
    log_dN = make_log_dN(a1, a2, mb, mu_q, sigma_q)
    m1 -> begin
        if m1 <= mlow
            zero(m1)
        elseif m1 > mhigh
            zero(m1)
        else
            n = Int(round(100*(log(m1)-log(mlow)))) + 2
            ms = exp.(range(log(mlow), stop=log(m1), length=n))
            dN = exp.(map(m2->log_dN(m1,m2), ms))
            trapz(ms, dN)
        end
    end
end

"""
    make_dNdm2(a1, a2, mb, beta)

Returns a function that computes the marginal density (unnormalized) in `m2`.
"""
function make_dNdm2(a1, a2, mb, mu_q, sigma_q)
    log_dN = make_log_dN(a1, a2, mb, mu_q, sigma_q)
    m2 -> begin
        if m2 < mlow
            zero(m2)
        elseif m2 > mhigh
            zero(m2)
        else
            n = Int(round(100*(log(mhigh) - log(m2)))) + 2
            ms = exp.(range(log(m2), stop=log(mhigh), length=n))
            dN = exp.(map(m1->log_dN(m1, m2), ms))
            trapz(ms, dN)
        end
    end
end

"""
Returns a function that computes the marginal (unnormalized) density in `q`.
"""
function make_dNdq(a1, a2, mb, mu_q, sigma_q)
    log_dN = make_log_dN(a1, a2, mb, mu_q, sigma_q)
    q -> begin
        if q < mlow/mhigh
            zero(q)
        elseif q > 1
            zero(q)
        else
            ml = mlow/q
            mh = mhigh
            n = Int(round(100*(log(mh)-log(ml)))) + 2
            ms = exp.(range(log(ml), stop=log(mh), length=n))
            dN = exp.(map(m1->log_dN(m1, q*m1) + log(m1), ms))
            trapz(ms, dN)
        end
    end
end

"""
    broken_pl_model(m1s, m2s, log_wts, m1s_sel, m2s_sel, log_pdraw, Ndraw)

Returns a Turing model for our broken power law mass function.

# Arguments

- `m1s`: An iterable containing iterables of the posterior samples for the
  source-frame primary mass for our observations.
- `m2s`: Same, but for secondary mass.
- `log_wts`: Iterable with iterables of the log of the prior weight assigned to
  each sample.
- `m1s_sel`: Samples of primary mass from the detected injections used to
  estimate detector sensitivity.
- `m2s_sel`: Secondary mass.
- `log_pdraw`: The log of the (normalized) density from which the injections
  were drawn.
- `Ndraw`: The number of injections drawn.

# Parameters

- `a1`: The low-mass power law slope.
- `a2`: The high-mass power law slope.
- `mb`: The break mass.
- `beta`: The pairing power law exponent.
- `R`: (generated) The volumetric merger rate per natural log mass squared at
  `m1 = m2 = mb`.
- `Neff_sel`: The effective number of samples (see [Farr
  (2019)](https://iopscience.iop.org/article/10.3847/2515-5172/ab1d5f)) in the
  Monte-Carlo selection integral.
"""
@model function broken_pl_model(m1s, m2s, log_wts, m1s_sel, m2s_sel, log_pdraw, Ndraw)
    Nobs = length(m1s)
    a1 ~ Uniform(0, 10)
    a2 ~ Uniform(-10, 0)
    mb ~ Uniform(sqrt(30), sqrt(200)) # Peaks somewhere around 10.
    mu_q ~ Uniform(0.2, 1.0)
    sigma_q ~ Uniform(0.1, 2)

    log_dN = make_log_dN(a1, a2, mb, mu_q, sigma_q)

    log_wts = map(m1s, m2s, log_wts) do mm1, mm2, log_wwt
        map(mm1, mm2, log_wwt) do m1, m2, log_wt
            log_dN(m1, m2) - log_wt
        end
    end

    Neff_samps = map(log_wts) do lws
        exp(2*logsumexp(lws) - logsumexp(2 .* lws))
    end

    log_likes = map(log_wts) do lw
        logsumexp(lw) - log(length(lw))
    end

    Turing.@addlogprob! sum(log_likes)

    log_wts_sel = map(m1s_sel, m2s_sel, log_pdraw) do m1, m2, logp
        log_dN(m1, m2) - logp
    end
    log_mu = logsumexp(log_wts_sel) - log(Ndraw)
    log_s2 = logsubexp(logsumexp(2 .* log_wts_sel) - 2.0*log(Ndraw), 2*log_mu - log(Ndraw))
    Neff_sel = exp(2*log_mu - log_s2)

    Turing.@addlogprob! -Nobs*log_mu

    mu_R = Nobs/exp(log_mu)
    sigma_R = sqrt(Nobs) / exp(log_mu)

    R = rand(Normal(mu_R, sigma_R))

    return (R = R, Neff_sel = Neff_sel, Neff_samps = Neff_samps)
end


using BackwardsLinalg
using LinearAlgebra:svd, Diagonal
using TensorOperations
using Flux, Test

function trg_svd(Ma, Dmax; tol::Float64=1e-12)
    U, S, V = svd(Ma)
    Dmax = min(searchsorted(S, tol, rev=true).stop, Dmax)
    D = isqrt(size(Ma, 1))
    #FS = view(S, 1:Dmax)  # this does not work!!!
    FS = S[1:Dmax]
    S1 = reshape(view(U,:,1:Dmax) .* sqrt.(FS'), (D, D, Dmax))
    S3 = reshape(sqrt.(FS) .* view(V',1:Dmax,:), (Dmax, D, D))
    S1, S3
end

function TRG(K::RT, Dcut::Int, no_iter::Int) where RT
    D = 2
    inds = 1:D

    M = [sqrt(cosh(K)) sqrt(sinh(K));
         sqrt(cosh(K)) -sqrt(sinh(K))]

    T = [mapreduce(a->M[a, i] * M[a, j] * M[a, k] * M[a, l], +, inds) for i in inds, j in inds, k in inds, l in inds]
    eltype(T) <: Tracker.TrackedReal && (T = Tracker.collect(T))

    lnZ = zero(RT)
    for n in 1:no_iter
        maxval = maximum(T)
        T = T/maxval
        lnZ += 2^(no_iter-n+1)*log(maxval)

        D = size(T, 1)

        Ma = reshape(permutedims(T, (3, 2, 1, 4)),  (D^2, D^2))
        Mb = reshape(permutedims(T, (4, 3, 2, 1)),  (D^2, D^2))

        S1, S3 = trg_svd(Ma, Dcut)
        S2, S4 = trg_svd(Mb, Dcut)

        # @tensoropt is much faster than @tensor
        @tensoropt T_new[r, u, l, d] := S1[w, a, r] * S2[a, b, u] * S3[l, b, g] * S4[d, g, w]

        T = T_new
    end
    trace = zero(RT)
    for i in 1:size(T, 1)
        trace += T[i, i, i, i]
    end
    lnZ += log(trace)
end

@testset "trg bp" begin
    Dcut = 24
    n = 3

    K = param(0.5)
    Z = TRG(K, Dcut, n)

    Tracker.back!(Z)
    res = Tracker.grad(K)

    δ = 1e-3
    Z0 = TRG(Tracker.data(K)-δ/2, Dcut, n)
    Z1 = TRG(Tracker.data(K)+δ/2, Dcut, n)
    ngrad = (Z1 - Z0)/δ
    @test isapprox(res, ngrad, atol=1e-2)
end


module Commons
  function backsolve(X,Y)

    row = size(X)[2]
    betacol = zeros(eltype(Y),row,1)

    iterator =row
    while iterator>0

      betacol[iterator] = (Y[iterator]- (X[iterator,:]*betacol)[1])/X[iterator,iterator]

      iterator-=1
    end

    return betacol

  end

end


#=
tests for PolyUtils module
=#

using Test
import Random
using LinearAlgebra
import Hypatia.PolyUtils
import Hypatia.PolyUtils.interpolate

function test_interp_domain(T::Type{<:Real})
    Random.seed!(1)
    @testset "domain interpolation n = $n, halfdeg = $halfdeg" for (n, halfdeg) in
        ((1, 5), (2, 2), (3, 1))
        free_dom = PolyUtils.FreeDomain{T}(n)
        box_dom = PolyUtils.BoxDomain{T}(-ones(T, n), ones(T, n))
        ball_dom = PolyUtils.BallDomain{T}(ones(T, n), one(T))
        Q = rand(T, n, n)
        Q = Symmetric(Q * Q')
        ellip_dom = PolyUtils.EllipsoidDomain{T}(ones(T, n), Q)

        for sample in (true, false)
            interp_options = (; sample = sample, sample_factor = 10)

            free = interpolate(free_dom, halfdeg; interp_options...)
            box = interpolate(box_dom, halfdeg; interp_options...)
            @test length(free.Ps) == 1
            @test length(box.Ps) == 1 + n
            @test free.U == box.U
            @test size(box.pts) == size(free.pts)
            @test size(box.Ps[1]) == size(free.Ps[1])
            @test norm(box.Ps[1]) ≈ norm(free.Ps[1]) atol=1e-1 rtol=1e-1

            if sample
                ball = interpolate(ball_dom, halfdeg; interp_options...)
                ellip = interpolate(ellip_dom, halfdeg; interp_options...)
                @test length(ball.Ps) == 2
                @test length(ellip.Ps) == 2
                @test free.U == ball.U == ellip.U
                @test size(ball.pts) == size(ellip.pts)
            end
        end
    end
end

function test_complex_interp(T::Type{<:Real})
    @testset "complex interp n = $n, halfdeg = $halfdeg, use_qr = $use_qr" for
        (n, halfdeg) in ((1, 3), (2, 2), (3, 1)), use_qr in (false, true)
        gs = [z -> 1 - sum(abs2, z)]
        (pts, Ps) = interpolate(Complex{T}, halfdeg, n, gs, [1], use_qr = use_qr)
        @test length(Ps) == 2
        L = binomial(n + halfdeg, n)
        U = L^2
        @test length(pts) == U
        @test size(Ps[1]) == (U, L)
    end
end

function test_quadrature(T::Type{<:Real})
    # univariate
    @testset "quadrature n = 1, halfdeg = $halfdeg" for halfdeg in 1:4
        free_dom = PolyUtils.FreeDomain{T}(1)
        free = interpolate(free_dom, halfdeg, sample = false, get_quadr = true)
        pts_polys = [sum(free.pts[i, 1] ^ d for d in
            0:(2 * halfdeg)) for i in 1:free.U]
        @test dot(free.w, pts_polys) ≈ sum(2 / (2i + 1) for i in 0:halfdeg)
    end

    # multivariate box
    @testset "quadrature n = $n, halfdeg = $halfdeg" for (n, halfdeg) in
        ((1, 5), (2, 2), (3, 1))
        box_dom = PolyUtils.BoxDomain{T}(fill(-1, n), fill(3, n))
        box = interpolate(box_dom, halfdeg, sample = false, get_quadr = true)
        @test sum(box.w) ≈ 2^n
    end
end


"""
    Addition of Two  normal random variables assuming
    zero covariance
"""
+(A::MvNormal, B::MvNormal) = MvNormal(a.μ + b.μ, A.Σ + B.Σ)
-(A::MvNormal, B::MvNormal) = MvNormal(a.μ - b.μ, A.Σ + B.Σ)

chol(A::MvNormal) = A.Σ.chol


#function spectral_line(x::AbstractArray; center::Real=0.0, width::Real=0.0, depth::Real=1.0, max_dist::Real=5.0 )
#   return spectral_line.(x,center=center,width=width,depth=depth, max_dist=max_dist)
#end

function spectral_line(x::Real; center::Real=0.0, width::Real=0.0, depth::Real=1.0, max_dist::Real=5.0 )
   return (x-center)/width > max_dist ? 1.0 : 1.0-depth*exp(-0.5*((x-center)/width)^2)
end

function spectral_line(x::Real; center::AbstractArray, width::AbstractArray, depth::AbstractArray, max_dist::Real=5.0 )
   @assert length(center) == length(width) == length(depth)
   output = one(typeof(x))
   for i in 1:length(center)
      output *= spectral_line(x,center=center[i],width=width[i],depth=depth[i], max_dist=max_dist)
   end
   return output
end

function make_toy_data(obs_size::Tuple{Int64,Int64},num_obs::Int = 0)
   instrument = :HPF
   data = EchelleObservationSet(obs_size, num_obs, instrument=instrument)
   lambda_min = 4000.0
   lambda_max = 8000.0
   num_lines = 1000
   line_widths = fill(5.0,num_lines)
   line_depths = fill(0.5,num_lines)
   line_centers = lambda_min.+(lambda_max-lambda_min).*rand(num_lines)

   snr = 100
   (num_obs, num_orders) = obs_size
   for i in 1:num_orders
      data.lambda[:,i,:] .= range(lambda_min+(lambda_max-lambda_min)*(i-1)/num_orders,
         stop=lambda_min+(lambda_max-lambda_min)*(i)/num_orders,length=obs_size[1])
   end
   data.flux .= spectral_line.(data.lambda, center=line_centers, width=line_widths, depth=line_depths, max_dist=5.0)
   data.var .= 1.0/snr
   data.flux .*= 1.0 .+ randn(size(data.flux))./snr
   return data
end


include("Linear.jl")
include("jacobian.jl")
using Base.Test

X = map(Float32, randn(5, 3))
l = Linear{Float32}(3, 4)

# Test grad. wrt. parameter W
Dfwd = jacobian_fwd(l, X, l.w)
Dbwd = jacobian_bwd(l, X, l.w, l.∂w)

@test_approx_eq Dfwd Dbwd

# Test grad. wrt. parameter b
Dfwd = jacobian_fwd(l, X, l.b)
Dbwd = jacobian_bwd(l, X, l.b, l.∂b)

@test_approx_eq Dfwd Dbwd

# Test grad. wrt. input X
Dfwd = jacobian_fwd(l, X)
Dbwd = jacobian_bwd(l, X)

@test_approx_eq Dfwd Dbwd


function reach_ASB07!(R::Vector{<:ReachSet},
                      Ω0::LazySet,
                      U::Union{ConstantInput, Nothing},
                      Φ::AbstractMatrix,
                      N::Int,
                      δ::Float64,
                      max_order::Int)
    # initial reach set
    t0, t1 = zero(δ), δ
    R[1] = ReachSet(Ω0, t0, t1)

    k = 2
    while k <= N
        Rₖ = overapproximate(Φ * R[k-1].X, Zonotope)
        if U != nothing
            Rₖ = minkowski_sum(Rₖ, next_set(U))
        end
        Rₖ = reduce_order(Rₖ, max_order)  # reduce order

        # store reach set
        t0 = t1
        t1 += δ
        if ! (genmat(Rₖ) isa Matrix)
            Rₖ = Zonotope(center(Rₖ), Matrix(genmat(Rₖ)))
        end
        R[k] = ReachSet(Rₖ, t0, t1)

        k += 1
    end
    return R
end


# An implementation of a FISTA-like method, where the smooth part of the objective function can be strongly convex.

using Base.Iterators
using ProximalAlgorithms.IterationTools
using ProximalCore: Zero
using LinearAlgebra
using Printf

"""
    SFISTAIteration(; <keyword-arguments>)

Iterator implementing the FISTA-like algorithm in [3].

This iterator solves strongly convex composite optimization problems of the form

    minimize f(x) + g(x),

where g is proper closed convex and f is a continuously differentiable function that is `mf`-strongly convex and whose gradient is
`Lf`-Lipschitz continuous.

The scheme is based on Nesterov's accelerated gradient method [1, Eq. (4.9)] and Beck's method for the convex case [2]. Its full
definition is given in [3, Algorithm 2.2.2.], and some analyses of this method are given in [3, 4, 5]. Another perspective is that
it is a special instance of [4, Algorithm 1] in which μh=0.

See also: [`SFISTA`](@ref).

# Arguments
- `x0`: initial point; must be in the domain of g.
- `f=Zero()`: smooth objective term.
- `g=Zero()`: proximable objective term.
- `mf=0` : strong convexity constant of f (see above).
- `Lf` : Lipschitz constant of ∇f (see above).

# References
1. Nesterov, Y. (2013). Gradient methods for minimizing composite functions. Mathematical Programming, 140(1), 125-161.
2. Beck, A., & Teboulle, M. (2009). A fast iterative shrinkage-thresholding algorithm for linear inverse problems. SIAM journal on imaging sciences, 2(1), 183-202.
3. Kong, W. (2021). Accelerated Inexact First-Order Methods for Solving Nonconvex Composite Optimization Problems. arXiv preprint arXiv:2104.09685.
4. Kong, W., Melo, J. G., & Monteiro, R. D. (2021). FISTA and Extensions - Review and New Insights. arXiv preprint arXiv:2107.01267.
5. Florea, M. I. (2018). Constructing Accelerated Algorithms for Large-scale Optimization-Framework, Algorithms, and Applications.
"""
Base.@kwdef struct SFISTAIteration{R,C<:Union{R,Complex{R}},Tx<:AbstractArray{C},Tf,Th}
    x0::Tx
    f::Tf = Zero()
    g::Th = Zero()
    Lf::R
    mf::R = real(eltype(Lf))(0.0)
end

Base.IteratorSize(::Type{<:SFISTAIteration}) = Base.IsInfinite()

Base.@kwdef mutable struct SFISTAState{R, Tx}
    λ::R                                # stepsize.
    yPrev::Tx                           # previous main iterate.
    y:: Tx = zero(yPrev)                # main iterate.
    xPrev::Tx = copy(yPrev)             # previous auxiliary iterate.
    x::Tx  = zero(yPrev)                # auxiliary iterate (see [3]).
    xt::Tx = zero(yPrev)                # prox center used to generate main iterate y.
    τ::R = real(eltype(yPrev))(1.0)     # helper variable (see [3]).
    a::R = real(eltype(yPrev))(0.0)     # helper variable (see [3]).
    APrev::R = real(eltype(yPrev))(1.0) # previous A (helper variable).
    A::R = real(eltype(yPrev))(0.0)     # helper variable (see [3]).
    gradf_xt::Tx = zero(yPrev)          # array containing ∇f(xt).
end

function Base.iterate(
    iter::SFISTAIteration,
    state::SFISTAState = SFISTAState(λ=1/iter.Lf, yPrev=copy(iter.x0))
)
    # Set up helper variables.
    state.τ = state.λ * (1 + iter.mf * state.APrev)
    state.a = (state.τ + sqrt(state.τ ^ 2 + 4 * state.τ * state.APrev)) / 2
    state.A = state.APrev + state.a
    state.xt .= (state.APrev / state.A) .* state.yPrev + (state.a / state.A) .* state.xPrev
    gradient!(state.gradf_xt, iter.f, state.xt)
    λ2 = state.λ / (1 + state.λ * iter.mf)
    # FISTA acceleration steps.
    prox!(state.y, iter.g, state.xt - λ2 * state.gradf_xt, λ2)
    state.x .= state.xPrev .+ (state.a / (1 + state.A * iter.mf)) .* ((state.y .- state.xt) ./ state.λ .+ iter.mf .* (state.y .- state.xPrev))
    # Update state variables.
    state.yPrev .= state.y
    state.xPrev .= state.x
    state.APrev = state.A
    return state, state
end

# Different stopping conditions (sc). Returns the current residual value and whether or not a stopping condition holds.
function check_sc(state::SFISTAState, iter::SFISTAIteration, tol, termination_type)
    if termination_type == "AIPP"
        # AIPP-style termination [4]. The main inclusion is: r ∈ ∂_η(f + h)(y).
        r = (iter.y0 - state.x) / state.A
        η = (norm(iter.y0 - state.y) ^ 2 - norm(state.x - state.y) ^ 2) / (2 * state.A)
        res = (norm(r) ^ 2 + max(η, 0.0)) / max(norm(iter.y0 - state.y + r) ^ 2, 1e-16)
    else
        # Classic (approximate) first-order stationary point [4]. The main inclusion is: r ∈ ∇f(y) + ∂h(y).
        λ2 = state.λ / (1 + state.λ * iter.mf)
        gradf_y, = gradient(iter.f, state.y)
        r = gradf_y - state.gradf_xt + (state.xt - state.y) / λ2
        res = norm(r)
    end
    return res, (res <= tol || res ≈ tol)
end

default_solution(::SFISTAIteration, state::SFISTAState) = state.y

"""
    SFISTA(; <keyword-arguments>)

Constructs the the FISTA-like algorithm in [3].

This algorithm solves strongly convex composite optimization problems of the form

    minimize f(x) + h(x),

where h is proper closed convex and f is a continuously differentiable function that is μ-strongly convex and whose gradient is
Lf-Lipschitz continuous.

The scheme is based on Nesterov's accelerated gradient method [1, Eq. (4.9)] and Beck's method for the convex case [2]. Its full
definition is given in [3, Algorithm 2.2.2.], and some analyses of this method are given in [3, 4, 5]. Another perspective is that
it is a special instance of [4, Algorithm 1] in which μh=0.

The returned object has type `IterativeAlgorithm{SFISTAIteration}`,
and can be called with the problem's arguments to trigger its solution.

See also: [`SFISTAIteration`](@ref), [`IterativeAlgorithm`](@ref).

# Arguments
- `maxit::Int=10_000`: maximum number of iteration
- `tol::1e-6`: tolerance for the default stopping criterion
- `stop::Function`: termination condition, `stop(::T, state)` should return `true` when to stop the iteration
- `solution::Function`: solution mapping, `solution(::T, state)` should return the identified solution
- `verbose::Bool=false`: whether the algorithm state should be displayed
- `freq::Int=100`: every how many iterations to display the algorithm state
- `display::Function`: display function, `display(::Int, ::T, state)` should display a summary of the iteration state
- `kwargs...`: additional keyword arguments to pass on to the `SFISTAIteration` constructor upon call

# References
1. Nesterov, Y. (2013). Gradient methods for minimizing composite functions. Mathematical Programming, 140(1), 125-161.
2. Beck, A., & Teboulle, M. (2009). A fast iterative shrinkage-thresholding algorithm for linear inverse problems. SIAM journal on imaging sciences, 2(1), 183-202.
3. Kong, W. (2021). Accelerated Inexact First-Order Methods for Solving Nonconvex Composite Optimization Problems. arXiv preprint arXiv:2104.09685.
4. Kong, W., Melo, J. G., & Monteiro, R. D. (2021). FISTA and Extensions - Review and New Insights. arXiv preprint arXiv:2107.01267.
5. Florea, M. I. (2018). Constructing Accelerated Algorithms for Large-scale Optimization-Framework, Algorithms, and Applications.
"""
SFISTA(;
    maxit=10_000,
    tol=1e-6,
    termination_type="",
    stop=(iter, state) -> check_sc(state, iter, tol, termination_type)[2],
    solution=default_solution,
    verbose=false,
    freq=100,
    display=(it, iter, state) -> @printf("%5d | %.3e\n", it, check_sc(state, iter, tol, termination_type)[1]),
    kwargs...
) = IterativeAlgorithm(SFISTAIteration; maxit, stop, solution, verbose, freq, display, kwargs...)


##### filter design routines.

function designfilter(  config::MiniMaxConfigType,
                        L::Int,
                        candidate_multiple::Int,
                        wfunc::Function,
                        dfunc::Function,
                        𝓧_selection_string::String,
                        fig_num::Int)

#
    ### prepare initial reference positions.
    𝓧 = Vector{BigFloat}(undef,0)
    if 𝓧_selection_string == "Chebyshev"
        𝓧 = collect( Chebyshev2ndnode(k,L+1) for k = 0:L+1 )

    elseif 𝓧_selection_string == "maxvol via QR"
        𝓧, 𝑖_unused = getinitialpositions(L, candidate_multiple)

    else
        println("Unknown 𝓧 selection string. Using Chebyshev.")
        𝓧 = collect( Chebyshev2ndnode(k,L+1) for k = 0:L+1 )
    end
    @assert length(unique(𝓧)) == length(𝓧) # sanity check

    ### run mini-max.
    𝓧_new, h_history, iters_ran, status_msg, fig_num = fitminimax(  wfunc,
                                                                dfunc,
                                                                𝓧,
                                                                fig_num,
                                                                config)

    ### get filter coefficients.

    # pose the solution as a Lagrange interpolated function.

    return 𝓧_new, h_history, iters_ran, status_msg, fig_num
end


using EchelleCCFs
using Test
import Statistics: mean

@testset "Check CCF accuracy" begin  # TODO repeat for multiple mask shapes
    @testset "Tophat mask w/ 1 line" begin
        Δv = 1000
        λ1 = 5025.0
        d1 = 0.5
        σ_v = 5000.0
        λlo = 5000
        λhi = 5050
        npix = 2000
        @test_nowarn BasicLineList([λ1], [d1])
        ll = BasicLineList([λ1], [d1])
        ll = BasicLineList([λ1], [d1])
        @test_nowarn BasicCCFPlan(line_list=ll, mask_shape=TopHatCCFMask(Δv))
        ccfpl = BasicCCFPlan(line_list=ll, mask_shape=TopHatCCFMask(Δv))
        @test_nowarn calc_ccf_v_grid(ccfpl)
        v_grid = calc_ccf_v_grid(ccfpl)
        @test length(v_grid) == calc_length_ccf_v_grid(ccfpl)
        λs = exp.(range(log(λlo),stop=log(λhi),length=npix))
        flux = ones(size(λs))
        flux .*= 1 .- d1.*exp.(-0.5.*((λs.-λ1)./λ1.*(EchelleCCFs.speed_of_light_mps./σ_v)).^2)
        @test_nowarn ccf_1D(λs, flux, ccfpl)
        ccf = ccf_1D(λs, flux, ccfpl)
        @test sum(ccf) > 0
        @test_nowarn  MeasureRvFromCCFGaussian()
        mrfcg = MeasureRvFromCCFGaussian(frac_of_width_to_fit=0.75)
        @test_nowarn measure_rv_from_ccf(v_grid,ccf, alg=mrfcg )
        vfit = measure_rv_from_ccf(v_grid,ccf, alg=mrfcg )
        @test vfit.rv ≈ 0  atol = 0.1   # TODO: tune tolerance better
        var = flux./10^5
        @test_nowarn ccf_1D(λs, flux, var, ccfpl)
        (ccf, ccf_var) = ccf_1D(λs, flux, var, ccfpl)
        @test sum(ccf) > 0
        @test_nowarn measure_rv_from_ccf(v_grid,ccf, ccf_var, alg=mrfcg )
        vfit = measure_rv_from_ccf(v_grid,ccf,ccf_var, alg=mrfcg )  # TODO: Check Why is this doing so much worse?
        #=
        using Plots
        global plt = plot()
        plot!(plt,v_grid,ccf,label=:none)
        scatter!(plt,v_grid,ccf,color=1,label=:none)
        plot!(plt,v_grid,sqrt.(ccf_var),label=:none)
        =#
        @test vfit.rv ≈ 0  atol = 0.5   # TODO: tune tolerance better
    end

    @testset "Tophat mask w/ 2 lines" begin
        Δv = 1000
        λ1 = 5025.0
        λ2 = 5075.0
        d1 = 0.5
        d2 = 0.25
        σ_v = 5000.0
        λlo = 5000
        λhi = 5100
        npix = 8000
        @test_nowarn BasicLineList([λ1,λ2], [d1,d2])
        ll = BasicLineList([λ1,λ2], [d1,d2])
        ll = BasicLineList([λ1,λ2], [d1,d2])
        @test_nowarn BasicCCFPlan(line_list=ll, mask_shape=TopHatCCFMask(Δv))
        ccfpl = BasicCCFPlan(line_list=ll, mask_shape=TopHatCCFMask(Δv))
        @test_nowarn calc_ccf_v_grid(ccfpl)
        v_grid = calc_ccf_v_grid(ccfpl)
        @test length(v_grid) == calc_length_ccf_v_grid(ccfpl)
        λs = exp.(range(log(λlo),stop=log(λhi),length=npix))
        flux = ones(size(λs))
        flux .*= 1 .- d1.*exp.(-0.5.*((λs.-λ1)./λ1.*(EchelleCCFs.speed_of_light_mps./σ_v)).^2)
        flux .*= 1 .- d2.*exp.(-0.5.*((λs.-λ2)./λ2.*(EchelleCCFs.speed_of_light_mps./σ_v)).^2)
        @test_nowarn ccf_1D(λs, flux, ccfpl)
        var = flux./10^5
        @test_nowarn ccf_1D(λs, flux, var, ccfpl)
        (ccf, ccf_var) = ccf_1D(λs, flux, var, ccfpl)

        @test_nowarn  MeasureRvFromCCFGaussian()
        mrfcg = MeasureRvFromCCFGaussian()
        @test_nowarn measure_rv_from_ccf(v_grid,ccf, alg=mrfcg )
        @test_nowarn measure_rv_from_ccf(v_grid,ccf, ccf_var, alg=mrfcg )
        vfit = measure_rv_from_ccf(v_grid,ccf, alg=mrfcg )
        @test vfit.rv ≈ 0  atol = 0.3   # TODO: tune tolerance better
        vfit = measure_rv_from_ccf(v_grid,ccf,ccf_var, alg=mrfcg )  # TODO: Check Why is this doing so much worse?
        #=
        using Plots
        global plt = plot()
        plot!(plt,v_grid,ccf,label=:none)
        scatter!(plt,v_grid,ccf,color=1,label=:none)
        plot!(plt,v_grid,sqrt.(ccf_var),label=:none)
        =#
        @test vfit.rv ≈ 0  atol = 0.3   # TODO: tune tolerance better

    end
    @testset "Gaussian mask" begin
        Δv = 7e3
        λ1 = 5025.0
        λ2 = 5075.0
        d1 = 0.5
        d2 = 0.25
        σ_v = 5000.0
        λlo = 5000
        λhi = 5100
        npix = 4000
        @test_nowarn BasicLineList([λ1,λ2], [d1,d2])
        ll = BasicLineList([λ1,λ2], [d1,d2])
        @test_nowarn BasicCCFPlan(line_list=ll, mask_shape=GaussianCCFMask(σ_v))
        ccfpl = BasicCCFPlan(line_list=ll, mask_shape=GaussianCCFMask(σ_v))
        @test_nowarn calc_ccf_v_grid(ccfpl)
        v_grid = calc_ccf_v_grid(ccfpl)
        @test length(v_grid) == calc_length_ccf_v_grid(ccfpl)
        λs = exp.(range(log(λlo),stop=log(λhi),length=npix))
        flux = ones(size(λs))
        flux .*= 1 .- d1.*exp.(-0.5.*((λs.-λ1)./λ1.*(EchelleCCFs.speed_of_light_mps./σ_v)).^2)
        flux .*= 1 .- d2.*exp.(-0.5.*((λs.-λ2)./λ2.*(EchelleCCFs.speed_of_light_mps./σ_v)).^2)
        @test_nowarn ccf_1D(λs, flux, ccfpl)
        var = flux./10^5
        @test_nowarn ccf_1D(λs, flux, var, ccfpl)
        (ccf, ccf_var) = ccf_1D(λs, flux, var, ccfpl)

        @test_nowarn measure_rv_from_ccf(v_grid,ccf, alg=MeasureRvFromCCFQuadratic() )
        @test_nowarn measure_rv_from_ccf(v_grid,ccf,ccf_var, alg=MeasureRvFromCCFQuadratic() )
        vfit = measure_rv_from_ccf(v_grid,ccf, alg=MeasureRvFromCCFQuadratic() )    # TODO: Check Why is this doing so much worse?
        @test vfit.rv ≈ 0  atol = 0.3   # TODO: tune tolerance better
        vfit = measure_rv_from_ccf(v_grid,ccf,ccf_var, alg=MeasureRvFromCCFQuadratic() )
        @test vfit.rv ≈ 0  atol = 0.3   # TODO: tune tolerance better

        @test_nowarn measure_rv_from_ccf(v_grid,ccf, alg=MeasureRvFromMinCCF() )
        vfit = measure_rv_from_ccf(v_grid,ccf, alg=MeasureRvFromMinCCF())
        @test vfit.rv ≈ 0  atol = 0.1   # TODO: tune tolerance better
        vfit = measure_rv_from_ccf(v_grid,ccf,ccf_var, alg=MeasureRvFromMinCCF())
        @test vfit.rv ≈ 0  atol = 0.5   # TODO: tune tolerance better

        #=
        @test_nowarn measure_rv_from_ccf(v_grid,ccf, alg=MeasureRvFromCCFCentroid() )
        vfit = measure_rv_from_ccf(v_grid,ccf)
        @test vfit.rv ≈ 0  atol = 500   # TODO: tune tolerance better
        =#
    end
    @testset "RV injection/extraction" begin
        # velocities
        amp = 1.0 # meters per second
        vel = amp .* sin.(range(0, 2π, length=100))
        mid = 5434.5
        dep = 0.75
        wid = 0.1
        mid_shift = mid .* EchelleCCFs.calc_doppler_factor.(vel)

        # compute gaussian line spectrum
        buffer = 0.75
        resolution = 7e5
        Δlnλ = (1.0 / resolution)
        wavs = exp.(range(log(mid - buffer), log(mid + buffer), step=Δlnλ))

        flux = zeros(length(wavs), length(vel))
        for i in eachindex(vel)
            flux[:,i] .= @. 1.0 - dep * exp(-((wavs-mid_shift[i])/wid)^2 / 2.0)
        end

        # measure ccf
        line_list = EchelleCCFs.BasicLineList([mid], [dep])
        speed_of_light = EchelleCCFs.speed_of_light_mps
        mask_width = speed_of_light/resolution
        mask_shape = TopHatCCFMask(mask_width)
        ccf_plan = BasicCCFPlan(line_list=line_list, mask_shape=mask_shape)
        v_grid = EchelleCCFs.calc_ccf_v_grid(ccf_plan)
        ccf = EchelleCCFs.ccf_1D(wavs, flux, ccf_plan)

        # measure velocities from CCF
        mrv_gauss = MeasureRvFromCCFGaussian(frac_of_width_to_fit=0.5)
        mrv_quad = MeasureRvFromCCFQuadratic(frac_of_width_to_fit=0.1)

        rvs_gauss = zeros(length(vel))
        sig_gauss = zeros(length(vel))
        rvs_quad = zeros(length(vel))
        sig_quad = zeros(length(vel))
        for i in eachindex(vel)
            rvs_gauss[i], sig_gauss[i] = mrv_gauss(v_grid, view(ccf, :, i))
            rvs_quad[i], sig_quad[i] = mrv_quad(v_grid, view(ccf, :, i))
        end
        err_gauss = abs.(rvs_gauss .- mean(rvs_gauss) .- vel)
        err_quad = abs.(rvs_quad .- mean(rvs_quad) .- vel)

        # test maximum errors are less than a cm/s
        @assert maximum(err_gauss) < 1e-2
        @assert maximum(err_quad) < 1e-2
    end
end


using BrainFlow

data = Array{Float64,3}(undef, 2, 4, 2)
data[:, :, 1] = [[6,3] [3,0] [1,5] [5,1]]
data[:, :, 2] = [[1,5] [5,1] [6,2] [2,2]]
data = permutedims(data, [3, 1, 2])
labels = [0.0, 1.0]

filters_true = [[-0.313406,-0.280803] [0.0792153, -0.480046]]
filters, eigvals = BrainFlow.get_csp(data, labels)

println("filters = ")
println(filters)

println("true filters =")
println(filters_true)


using Plots, MAT

nprocs = (2, 2) # nprocs (x, y) dim

@views function vizme2D_mpi(nprocs)
    T  = []
    ip = 1
    for ipx = 1:nprocs[1]
        for ipy = 1:nprocs[2]
            file = matopen("T_$(ip-1).mat"); T_loc = read(file, "T"); close(file)
            nx_i, ny_i = size(T_loc,1)-2, size(T_loc,2)-2
            ix1, iy1   = 1+(ipx-1)*nx_i, 1+(ipy-1)*ny_i
            if (ip==1)  T = zeros(nprocs[1]*nx_i, nprocs[2]*ny_i)  end
            T[ix1:ix1+nx_i-1,iy1:iy1+ny_i-1] .= T_loc[2:end-1,2:end-1]
            ip += 1
        end
    end
    display(heatmap(T', framestyle=:box, aspect_ratio=1, xlims=(1, size(T,1)), ylims=(1, size(T,2)), c=:hot, title="heat diffusion"))
    return
end

vizme2D_mpi(nprocs)


using Test
using AlgebraicDependencies

testnum = []
testdep = []

macro num(numbers...)
    push!(testnum, collect(numbers))
end

macro dep(deps...)
    push!(testdep, collect(deps))
end

gsym(n::Int) = [Symbol("v$i") for i in 1:n]

# ------------------------------------------------------------------------------

@num 1/2 2
@dep v1*v2-1

@num 1/2 1 2
@dep v2-1 v1*v3-1

# ------------------------------------------------------------------------------

for (n, d) in zip(testnum, testdep)
    @test dependencies(n; variables=gsym(length(n))) == d
end


@generated function convert(::Type{Monomial{v1}}, m::Monomial) where {v1}
    args = Any[0 for v in v1]
    v2 = variables(m)
    i1 = 1
    i2 = 1
    while i1 <= length(v1) && i2 <= length(v2)
        if v1[i1] == v2[i2]
            args[i1] = :(exponents(m)[$i2])
            i1 += 1
            i2 += 1
        else
            i1 += 1
        end
    end
    if i2 <= length(v2)
        throw(InexactError())
    end
    :(Monomial{v1, $(Expr(:tuple, args...))}())
end

convert(::Type{Monomial}, v::Variable) = Monomial{(v,), (1,)}()
convert(::Type{M}, v::Variable) where {M <: Monomial} = convert(M, Monomial(v))

convert(::Type{AbstractPolynomial}, terms::Tuple{Vararg{Term}}) = Polynomial(terms)

convert(::Type{Term}, m::Monomial) = Term(1, m)
convert(::Type{Term{T, M}}, m::Monomial) where {T, M} = Term(one(T), convert(M, m))
convert(::Type{Term}, v::Variable) = Term(1, Monomial(v))
convert(::Type{Term{T, M}}, v::Variable) where {T, M} = convert(Term{T, M}, Monomial(v))
convert(::Type{Term{T, M}}, x) where {T, M} = Term{T, M}(convert(T, x))


# Index manipulations
#---------------------
"""
    permute(tsrc::AbstractTensorMap{S}, p1::NTuple{N₁, Int}, p2::NTuple{N₂, Int} = ())
        -> tdst::TensorMap{S, N₁, N₂}

Permute the indices of `tsrc::AbstractTensorMap{S}` such that a new tensor
`tdst::TensorMap{S, N₁, N₂}` is obtained, with indices in `p1` playing the role of the
codomain or range of the map, and indices in `p2` indicating the domain.

To permute into an existing `tdst`, see [`add!`](@ref)
"""
function permute(t::TensorMap{S},
                    p1::IndexTuple{N₁},  p2::IndexTuple{N₂}=();
                    copy::Bool = false) where {S, N₁, N₂}
    cod = ProductSpace{S, N₁}(map(n->space(t, n), p1))
    dom = ProductSpace{S, N₂}(map(n->dual(space(t, n)), p2))
    # share data if possible
    if !copy
        if p1 === codomainind(t) && p2 === domainind(t)
            return t
        elseif has_shared_permute(t, p1, p2)
            return TensorMap(reshape(t.data, dim(cod), dim(dom)), cod, dom)
        end
    end
    # general case
    @inbounds begin
        return add!(true, t, false, similar(t, cod←dom), p1, p2)
    end
end

function permute(t::AdjointTensorMap{S}, p1::IndexTuple, p2::IndexTuple=();
                    copy::Bool = false) where {S}
    p1′ = map(n->adjointtensorindex(t, n), p2)
    p2′ = map(n->adjointtensorindex(t, n), p1)
    adjoint(permute(adjoint(t), p1′, p2′; copy = copy))
end

function has_shared_permute(t::TensorMap, p1, p2)
    if p1 === codomainind(t) && p2 === domainind(t)
        return true
    elseif sectortype(t) === Trivial
        stridet = i->stride(t[], i)
        sizet = i->size(t[], i)
        canfuse1, d1, s1 = TO._canfuse(sizet.(p1), stridet.(p1))
        canfuse2, d2, s2 = TO._canfuse(sizet.(p2), stridet.(p2))
        return canfuse1 && canfuse2 && s1 == 1 && (d2 == 1 || s2 == d1)
    else
        return false
    end
end

function has_shared_permute(t::AdjointTensorMap, p1, p2)
    p1′ = adjointtensorindices(t, p2)
    p2′ = adjointtensorindices(t, p1)
    return has_shared_permute(t', p1′, p2′)
end

"""
    permute!(tdst::AbstractTensorMap{S, N₁, N₂}, tsrc::AbstractTensorMap{S},
        p1::IndexTuple{N₁}, p2::IndexTuple{N₂}=()
    ) -> tdst

Permute the indices of `tsrc` and write the result into `tdst`, with indices in `p1` playing
the role of the codomain or range of the map `tdst`, and indices in `p2` indicating the
domain of `tdst`.

For more details see [`add!`](@ref), of which this is a special case.
"""
@propagate_inbounds Base.permute!(tdst::AbstractTensorMap{S, N₁, N₂},
                                    tsrc::AbstractTensorMap{S},
                                    p1::IndexTuple{N₁},
                                    p2::IndexTuple{N₂}=()) where {S, N₁, N₂} =
    add_permute!(true, tsrc, false, tdst, p1, p2)

# Braid
function braid(t::TensorMap{S}, levels::IndexTuple,
                    p1::IndexTuple, p2::IndexTuple=();
                    copy::Bool = false) where {S}
    @assert length(levels) == numind(t)
    if BraidingStyle(sectortype(S)) isa SymmetricBraiding
        return permute(t, p1, p2; copy = copy)
    end
    if !copy && p1 == codomainind(t) && p2 == domainind(t)
        return t
    end
    # general case
    cod = ProductSpace{S}(map(n->space(t, n), p1))
    dom = ProductSpace{S}(map(n->dual(space(t, n)), p2))
    @inbounds begin
        return add_braid!(true, t, false, similar(t, cod←dom), p1, p2, levels)
    end
end
@propagate_inbounds braid!(tdst::AbstractTensorMap{S, N₁, N₂},
                                tsrc::AbstractTensorMap{S},
                                levels::IndexTuple,
                                p1::IndexTuple{N₁},
                                p2::IndexTuple{N₂}=()) where {S, N₁, N₂} =
    add!(true, tsrc, false, tdst, p1, p2, levels)

# Transpose
LinearAlgebra.transpose!(tdst::AbstractTensorMap, tsrc::AbstractTensorMap,
                            p1::IndexTuple = reverse(domainind(tsrc)),
                            p2::IndexTuple = reverse(codomainind(tsrc))) =
    add_transpose!(true, tsrc, false, tdst, p1, p2)

function LinearAlgebra.transpose(t::TensorMap{S},
                                    p1::IndexTuple = reverse(domainind(t)),
                                    p2::IndexTuple = reverse(codomainind(t));
                                    copy::Bool = false) where {S}
    if sectortype(S) === Trivial
        return permute(t, p1, p2; copy = copy)
    end
    if !copy && p1 == codomainind(t) && p2 == domainind(t)
        return t
    end
    # general case
    cod = ProductSpace{S}(map(n->space(t, n), p1))
    dom = ProductSpace{S}(map(n->dual(space(t, n)), p2))
    @inbounds begin
        return add_transpose!(true, t, false, similar(t, cod←dom), p1, p2)
    end
end

function LinearAlgebra.transpose(t::AdjointTensorMap{S},
                                    p1::IndexTuple = reverse(domainind(t)),
                                    p2::IndexTuple = reverse(codomainind(t));
                                    copy::Bool = false) where {S}
    p1′ = map(n->adjointtensorindex(t, n), p2)
    p2′ = map(n->adjointtensorindex(t, n), p1)
    adjoint(transpose(adjoint(t), p1′, p2′; copy = copy))
end

# Twist
twist(t::AbstractTensorMap, i::Int; inv::Bool = false) = twist!(copy(t), i; inv = inv)

function twist!(t::AbstractTensorMap, i::Int; inv::Bool = false)
    if i > numind(t)
        msg = "Can't twist index $i of a tensor with only $(numind(t)) indices."
        throw(ArgumentError(msg))
    end
    BraidingStyle(sectortype(t)) == Bosonic() && return t
    N₁ = numout(t)
    for (f1, f2) in fusiontrees(t)
        θ = i <= N₁ ? twist(f1.uncoupled[i]) : twist(f2.uncoupled[i-N₁])
        inv && (θ = θ')
        rmul!(t[f1, f2], θ)
    end
    return t
end

# Fusing and splitting


export AbstractNorm, InfNorm, WeightedNorm, EuclideanNorm, distance, norm, init!, update!

using Base: @propagate_inbounds
import LinearAlgebra: norm

"""
    AbstractNorm

An `AbstractNorm` represents any norm of a vector space.
All norms are callable. `norm(x)` computes the norm of `x` and `norm(x,y)` computes the distance
`norm(x - y).`
"""
abstract type AbstractNorm end

"""
    distance(u, v, norm::AbstractNorm)

Compute the distance ||u-v|| with respect to the given norm `norm`.
"""
distance(u, v, norm::AbstractNorm) = MethodError(distance, (u, v, norm))

"""
    norm(u, norm::AbstractNorm)

Compute the norm ||u|| with respect to the given norm `norm`.
"""
LinearAlgebra.norm(u, norm::AbstractNorm) = MethodError(LinearAlgebra.norm, (u, norm))

(N::AbstractNorm)(x) = norm(x, N)
(N::AbstractNorm)(x, y) = distance(x, y, N)

##################
## WeightedNorm ##
##################

Base.@kwdef struct WeightedNormOptions
    scale_min::Float64 = 1e-4
    scale_abs_min::Float64 = 1e-6
    scale_max::Float64 = exp2(div(1023, 2))
end

Base.show(io::IO, opts::WeightedNormOptions) = print_fieldnames(io, opts)
Base.show(io::IO, ::MIME"application/prs.juno.inline", opts::WeightedNormOptions) = opts

"""
    WeightedNorm(d::AbstractVector, norm::AbstractNorm; options...)
    WeightedNorm(norm::AbstractNorm, n::Integer; options...)
    WeightedNorm(norm::AbstractNorm, x::AbstractVector; options...)

A `WeightedNorm` represents a weighted variant of norm `norm` of a `n`-dimensional vector space.`
A norm ``||x||`` is weighted by introducing a vector of additional weights `d` such that
the new norm is ``||D⁻¹x||`` where ``D`` is the diagonal matrix with diagonal ``d`.
The `WeightedNorm` is desigened to change the weights dynamically by using [`init!(::WeightedNorm, x)`](@ref)
and [`update!(::WeightedNorm, x)`](@ref). The weights are there constructed such that
``||D⁻¹x|| ≈ 1.0``.
The weights can be accessed and changed by indexing.

### Options

* `scale_min = sqrt(eps())`: The minimal size of `dᵢ` is `scale_min` time the (weighted) norm of `x`.
* `scale_abs_min = min(scale_min^2, 200 * sqrt(eps()))`: The absolute minimal size of `dᵢ`.
* `scale_max = 1.0 / eps() / sqrt(2)`: The absolute maximal size of `dᵢ`

"""
struct WeightedNorm{N<:AbstractNorm} <: AbstractNorm
    weights::Vector{Float64}
    norm::N
    options::WeightedNormOptions
end
function WeightedNorm(
    weights::Vector{Float64},
    norm::AbstractNorm,
    opts = WeightedNormOptions(),
)
    WeightedNorm(weights, norm, opts)
end
WeightedNorm(norm::AbstractNorm, x::AbstractVector) = WeightedNorm(norm, length(x))
WeightedNorm(norm::AbstractNorm, n::Integer) = WeightedNorm(ones(n), norm)

"""
    weights(WN::WeightedNorm)

Returns the weights of the weighted norm.
"""
weights(WN::WeightedNorm) = WN.weights

Base.length(WN::WeightedNorm) = length(WN.weights)
Base.size(WN::WeightedNorm) = (length(WN),)
Base.ndims(::Type{<:WeightedNorm}) = 1
Base.copyto!(WN::WeightedNorm, x) = copyto!(WN.weights, x)
@propagate_inbounds Base.getindex(WN::WeightedNorm, i::Integer) = getindex(WN.weights, i)
@propagate_inbounds Base.setindex!(WN::WeightedNorm, x, i::Integer) =
    setindex!(WN.weights, x, i)

"""
    init!(w::WeightedNorm, x::AbstractVector)

Setup the weighted norm `w` for `x`.
"""
function init!(w::WeightedNorm, x)
    point_norm = w.norm(x)
    for i = 1:length(w)
        wᵢ = fast_abs(x[i])
        if wᵢ < w.options.scale_min * point_norm
            wᵢ = w.options.scale_min * point_norm
        elseif wᵢ > w.options.scale_max * point_norm
            wᵢ = w.options.scale_max * point_norm
        end
        w[i] = max(wᵢ, w.options.scale_abs_min)
    end
    w
end
init!(n::AbstractNorm, ::Any) = n

"""
    update!(w::WeightedNorm, x)

Update the weighted norm `w` for `x`, this will interpolate between the previous weights
and the norm of `x`.
"""
function update!(w::WeightedNorm, x)
    norm_x = w(x)
    for i = 1:length(x)
        wᵢ = (fast_abs(x[i]) + w[i]) / 2
        if wᵢ < w.options.scale_min * norm_x
            wᵢ = w.options.scale_min * norm_x
        elseif wᵢ > w.options.scale_max * norm_x
            wᵢ = w.options.scale_max * norm_x
        end
        if isfinite(wᵢ)
            w[i] = max(wᵢ, w.options.scale_abs_min)
        end
    end
    w
end
update!(n::AbstractNorm, ::Any) = n

####################
## Specific norms ##
####################
"""
    InfNorm <: AbstractNorm

The infinity or [maximum norm](https://en.wikipedia.org/wiki/Norm_(mathematics)#Maximum_norm).
"""
struct InfNorm <: AbstractNorm end

function distance(x, y, ::InfNorm)
    n = length(x)
    @boundscheck n == length(y)
    @inbounds dmax = abs2(x[1] - y[1])
    for i = 2:n
        @inbounds dᵢ = abs2(x[i] - y[i])
        dmax = Base.FastMath.max_fast(dmax, dᵢ)
    end
    d = sqrt(dmax)
    if isinf(d)
        @inbounds dmax = abs(x[1] - y[1])
        for i = 2:n
            @inbounds dᵢ = abs2(x[i] - y[i])
            dmax = max(dmax, dᵢ)
        end
        return dmax
    else
        return d
    end
end
function distance(x, y, w::WeightedNorm{InfNorm})
    inf_distance(x, y, weights(w))
end
function inf_distance(x, y, w)
    n = length(x)
    @boundscheck n == length(w)
    @inbounds dmax = abs2((x[1] - y[1]) / w[1])
    for i = 2:n
        @inbounds dᵢ = abs2((x[i] - y[i]) / w[i])
        dmax = Base.FastMath.max_fast(dmax, dᵢ)
    end
    d = sqrt(dmax)
    if isinf(d)
        @inbounds dmax = abs((x[1] - y[1]) / w[1])
        for i = 2:n
            @inbounds dᵢ = abs((x[i] - y[i]) / w[i])
            dmax = max(dmax, dᵢ)
        end
        return dmax
    else
        return d
    end
end

function norm(x, ::InfNorm)
    n = length(x)
    @inbounds dmax = abs2(x[1])
    for i = 2:n
        @inbounds dᵢ = abs2(x[i])
        dmax = Base.FastMath.max_fast(dmax, dᵢ)
    end
    d = sqrt(dmax)

    if isinf(d)
        @inbounds dmax = abs(x[1])
        for i = 2:n
            @inbounds dᵢ = abs(x[i])
            dmax = max(dmax, dᵢ)
        end
        return dmax
    else
        return d
    end
end
function norm(x, w::WeightedNorm{InfNorm})
    n = length(x)
    @boundscheck n == length(w)
    @inbounds dmax = abs2(x[1] / w[1])
    for i = 2:n
        @inbounds dᵢ = abs2(x[i] / w[i])
        dmax = Base.FastMath.max_fast(dmax, dᵢ)
    end
    d = sqrt(dmax)

    if isinf(d)
        @inbounds dmax = abs(x[1] / w[1])
        for i = 2:n
            @inbounds dᵢ = abs(x[i] / w[i])
            dmax = max(dmax, dᵢ)
        end
        return dmax
    else
        return d
    end
end

"""
    EuclideanNorm <: AbstractNorm

The Euclidean resp. 2-norm.
"""
struct EuclideanNorm <: AbstractNorm end

@inline function distance(x, y, ::EuclideanNorm)
    n = length(x)
    @boundscheck n == length(y)
    @inbounds d = abs2(x[1] - y[1])
    for i = 2:n
        @inbounds d += abs2(x[i] - y[i])
    end
    sqrt(d)
end
@inline function norm(x, ::EuclideanNorm)
    n = length(x)
    @inbounds d = abs2(x[1])
    for i = 2:n
        @inbounds d += abs2(x[i])
    end
    sqrt(d)
end


"""
# Loc

a two-dimensional point (x, y). represents a location on a Quoridor
board

"""
immutable Loc
    x::Int
    y::Int
end

(+)(l1::Loc, l2::Loc) = Loc(l1.x + l2.x, l1.y + l2.y)
(-)(l1::Loc, l2::Loc) = Loc(l1.x - l2.x, l1.y - l2.y)

# TODO bounds check on these somewhere?
left(l::Loc) = l - Loc(1,0)
right(l::Loc) = l + Loc(1,0)
above(l::Loc) = l + Loc(0,1)
below(l::Loc) = l - Loc(0,1)

const DIRS = Vector4{Function}(above, below, left, right)
const VERTICAL = Vector2{Function}(above, below)
const HORIZONTAL = Vector2{Function}(left, right)

const PERPENDICULAR = Dict(left=>VERTICAL, right=>VERTICAL,
                           above=>HORIZONTAL, below=>HORIZONTAL)

isless(l1::Loc, l2::Loc) = isless(l1.x + l1.y, l2.x + l2.y)
mean(l1::Loc, l2::Loc) = Loc((l1.x + l2.x) / 2, (l1.y + l2.y) / 2)

convert(::Type{Loc}, tup::Tuple{Int, Int}) = Loc(tup[1], tup[2])

show(io::IO, loc::Loc) = write(io, "Loc($(loc.x), $(loc.y))")

isoutofbounds(loc::Loc) = (loc.x < 1 || loc.x > 9 ||
                           loc.y < 1 || loc.y > 9)

inbounds(loc::Loc) = !isoutofbounds(loc)

function adjacent(loc::Loc)
    Vector4{Loc}(above(loc), below(loc), right(loc), left(loc))
end


Base.promote_rule(::Type{Decimal}, ::Type{<:Real}) = Decimal

# override definitions in Base
Base.promote_rule(::Type{BigFloat}, ::Type{Decimal}) = Decimal
Base.promote_rule(::Type{BigInt}, ::Type{Decimal}) = Decimal

# Addition
# To add, convert both decimals to the same exponent.
# (If the exponents are different, use the smaller exponent
# to make sure we're adding integers.)
function +(x::Decimal, y::Decimal)
    cx = (-1)^x.s * x.c * 10^max(x.q - y.q, 0)
    cy = (-1)^y.s * y.c * 10^max(y.q - x.q, 0)
    s = (abs(cx) > abs(cy)) ? x.s : y.s
    c = BigInt(cx) + BigInt(cy)
    normalize(Decimal(s, abs(c), min(x.q, y.q)))
end

# Negation
-(x::Decimal) = Decimal((x.s == 1) ? 0 : 1, x.c, x.q)

# Subtraction
-(x::Decimal, y::Decimal) = +(x, -y)

# Multiplication
function *(x::Decimal, y::Decimal)
    s = (x.s == y.s) ? 0 : 1
    normalize(Decimal(s, BigInt(x.c) * BigInt(y.c), x.q + y.q))
end

# Inversion
function Base.inv(x::Decimal)
    str = string(x)
    if str[1] == '-'
        str = str[2:end]
    end
    b = ('.' in str) ? length(split(str, '.')[1]) : 0
    c = round(BigInt(10)^(-x.q + DIGITS) / x.c)
    q = (x.q < 0) ? 1 - b - DIGITS : -b - DIGITS
    normalize(Decimal(x.s, c, q))
end

# Division
/(x::Decimal, y::Decimal) = x * inv(y)

# Exponentiation, please test

function ^(x::Decimal, y::Decimal)
    #(-1.23)^(something) code to catch errors prior to the routine
    x.s == 1 && begin
        (y.q < 0  && throw(DomainError(y),"Exponentiation yielding a complex result requires a complex argument."))
    end
    y.s == 1 && (return ^(inv(x),-y))
    y.q < 0 && (return Decimal(BigFloat(x)^BigFloat(y))) #add correction to precision, its too much

    if iseven(y.c) #squared number
        zs = 0
    else
        if x.s == 0
            zs =0
        else
            zs = 1
        end
    end

        zc = x.c^((y.c)*(10^y.q))
        zq = (x.q)*((y.c)*(10^y.q))
        return  Decimal(zs,zc,zq)

end


#!/usr/bin/julia
# Looks for a bound.
# I'm not sure if it's currently working.

"""
Ternary search (based on Wikipedia). This searches for an integer
    value maximizing f.
  f: function to maximize
  left, right: bounds
  Returns: value maximizing f
"""
function ternarySearch(f, left, right)
  print(string(left) * " " * string(right) * "\n")
  # left and right are the current bounds;
  # the maximum is between them
  if abs(right - left) <= 1
    return trunc((left+right)/2)
  end
  leftThird = ((2*left + right)/3)
  rightThird = ((left + 2*right)/3)
  if f(trunc(leftThird)) < f(trunc(rightThird))
    return ternarySearch(f, leftThird, right)
  else
    return ternarySearch(f, left, rightThird)
  end
end

"""
Counting bound, based on Shannon's argument.
  m: number of edges
  w: number of 'wires' -- that is, log2(number of functions),
    which is the number of bits needed to specify a function
  Returns: average number of gates required to computea
    any of those functions. (This may not be an integer).
"""
function countingBound(m, w)
  b = m - 0.5
  # the "-1" here is because this is the average, not the max.
  sqrt(2*w + b*b) - b - 1
end

"""
  Number of maximal hypercliques of some size.
  Note that k < r < n .
  k: number of vertices per hyperedge
  r: number of vertices in the clique
  n: vertices in the larger graph
  Returns: number of maximal hypercliques
"""
function numMaximalCliques(k, r, n)
  k = BigInt(k)
  r = BigInt(r)
  n = BigInt(n)
  one = BigInt(1)
  two = BigInt(2)

  # expected number of r-cliques
  numRCliques = Rational(binomial(n, r),
    two ^ binomial(r, k))

  # probability that one of those is not covered by a larger clique
  a = two ^ binomial(r, k-one)
  p = Rational(a-one, a) ^ (n-r)
  # println(a)
  # println(p)
  # result is number of cliques, * prob. they're maximal
  numRCliques * p
end

"""
  Bound, using one size of clique.
    This is approximate, because it's approximating the number of
    edges and cliques.
  k: number of vertices per hyperedge
  r: number of vertices for the clique finder
  n: vertices in the larger graph
  Returns: number of NAND gates
"""
function bound2(k, r, n)
  k = BigInt(k)
  r = BigInt(r)
  n = BigInt(n)
  one = BigInt(1)
  two = BigInt(2)

  # the bound on the total number of gates
  minTotalGates = ceil(countingBound(binomial(k, two), binomial(n, k)))

  # expected number of edges "left over" (not covered
  # by an r-clique)
  # p = (1 - 2 ^ (-(choose(r, k)-1))) ^ choose(n-k, r-k)
  a = two ^ (binomial(r, k) - one)
  # p = Rational(a-one, a) ^ binomial(n-k, r-k)
  # print("0\n")
  b = binomial(n-k, r-k)
  p = Rational((a-one)^b, a^b)
  # print("1\n")
  # num.left.over = (p/2) * choose(n, k)
  numLeftOver = p * binomial(n, k) / two
  # print("2\n")

  # expected number of r-cliques
  numRCliques = Rational(binomial(BigInt(n), BigInt(r)),
    BigInt(2) ^ binomial(r, k))
  # print("3\n")
  # the remaining gates are part of some number of
  # clique detectors
  r = (minTotalGates - numLeftOver) / numRCliques
end

"""
  Bound, using one size of clique.
    This is approximate, as it's estimating the number of edges
    and cliques. It also approximates the "(1 - 1/a)^..." bit.
  k: number of vertices per hyperedge
  r: number of vertices for the clique finder
  n: vertices in the larger graph
  Returns: number of NAND gates
"""
function bound3(k, r, n)
  k = BigInt(k)
  r = BigInt(r)
  n = BigInt(n)
  one = BigInt(1)
  two = BigInt(2)

  # the bound on the total number of gates (in the larger graph)
  minTotalGates = ceil(countingBound(binomial(k, two), binomial(n, k)))

  # expected number of edges "left over" (not covered by an r-clique)
  # First, the probability of a hyperedge being "left over"
  # (not covered by an r-clique) is
  #
  # p = (1 - 2 ^ (-(choose(r, k)-1))) ^ choose(n-k, r-k)
  #
  # We use an approximate upper bound (which should be good
  # for large numbers).
  logP = - binomial(n-k, r-k) / (two ^ (binomial(r, k) - one))
  # (log of) number of edges left over
  logNumLeftOver = logP + log(binomial(n, k) / two)

  # expected number of r-cliques
  # (Using right-shift for division, and adding 1 in case
  # of rounding error)
  numRCliques = (binomial(n, r) >> binomial(r, k)) + 1
  # the remaining gates are part of some number of
  # clique detectors
  r = (minTotalGates - exp(logNumLeftOver)) / numRCliques
end

function printBound1(k, r, n)
  b = bound1(k,r,n)
  print(string(k, ",", r, ",", n, ",", b, "\n"))
end

# print("k,r,n,bound\n")
# for n in 100:20:600
#   printBound1(5, 6, n)
# end