using ComputedFieldTypes
using DifferentialForms
using Random
using SimplexPoly
using StaticArrays
using Test

# Ignore a statement
macro DISABLED(expr)
    return quote end
end

const Dmax = 6

# Set reproducible random number seed
Random.seed!(0)

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

# Random terms
const maxrandpower = 5
function Base.rand(rng::AbstractRNG,
                   ::Random.SamplerType{Term{D,T}}) where {D,T}
    return Term{D,T}(SVector{D,Int}(rand(rng, 0:maxrandpower, D)),
                     T(rand(rng, -10:10)))
end

@testset "Terms as vector space D=$D" for D in 0:Dmax
    T = Int

    for iter in 1:100
        x = rand(Term{D,T})
        n = map(c -> zero(T), x)
        y = map(c -> T(rand(-10:10)), x)
        z = map(c -> T(rand(-10:10)), x)
        a = T(rand(-10:10))
        b = T(rand(-10:10))

        @test n == n

        f(x) = 2 * x
        g(x) = x + 1
        h(x) = x^2
        @test map(identity, x) == x
        @test map(f ∘ identity, x) == map(f, x)
        @test map(identity ∘ f, x) == map(f, x)
        @test map(h ∘ g, map(f, x)) == map(h, map(g ∘ f, x))

        @test n + x == x
        @test x + n == x
        @test x + y == y + x
        @test x + (y + z) == (x + y) + z
        @test +x == x

        @test x + (-x) == n
        @test -(-x) == x
        @test -n == n
        @test x - y == x + (-y)

        @test zero(T) * x == n
        @test one(T) * x == x
        @test (-one(T)) * x == -x
        @test a * x == x * a
        @test (a * b) * x == a * (b * x)
        @test (a + b) * x == a * x + b * x
        @test a * (x + y) == a * x + a * y

        # Cannot test this with integers
        # if a != zero(T)
        #     @test (a * x) / a == x
        #     @test a \ (a * x) == x
        # end
    end
end

@testset "Terms as ring D=$D" for D in 0:Dmax
    T = Int

    for iter in 1:100
        n = zero(Term{D,T})
        e = one(Term{D,T})
        x = rand(Term{D,T})
        x′ = map(c -> T(rand(-10:10)), x)
        y = rand(Term{D,T})
        z = rand(Term{D,T})
        a = T(rand(-10:10))
        b = T(rand(-10:10))
        i = rand(0:10)
        j = rand(0:10)

        @test n == n
        @test e == e
        @test e != n
        @test iszero(n)
        @test isone(e)

        # This is not true for terms, only for polynomials
        # @test n * x == n
        # @test x * n == n
        @test iszero(n * x)
        @test iszero(x * n)
        @test e * x == x
        @test x * e == x
        @test x * y == y * x
        @test x * (y * z) == (x * y) * z
        @test (a * x) * y == a * (x * y)
        @test x * (y * a) == (x * y) * a
        @test (x + x′) * y == x * y + x′ * y

        @test x^0 == e
        @test x^1 == x
        @test x^(i + j) == x^i * x^j
    end
end

@testset "Derivatives of terms D=$D" for D in 1:Dmax
    T = Int

    for iter in 1:100
        x = rand(Term{D,T})
        n = map(c -> zero(T), x)
        y = map(c -> T(rand(-10:10)), x)
        a = T(rand(-10:10))

        dir = rand(1:D)

        dx = deriv(x, dir)
        if x.powers[dir] == 0
            @test dx.coeff == 0
        else
            @test dx.powers[dir] == x.powers[dir] - 1
        end

        @test deriv(x + y, dir) == deriv(x, dir) + deriv(y, dir)
        @test deriv(a * x, dir) == a * deriv(x, dir)

        # @test deriv(x * y, dir) == deriv(x, dir) * y + x * deriv(y, dir)

        κx = koszul(x, dir)
        @test κx.powers[dir] == x.powers[dir] + 1

        @test koszul(x + y, dir) == koszul(x, dir) + koszul(y, dir)
        @test koszul(a * x, dir) == a * koszul(x, dir)
    end
end

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

# Random polynomials
function Base.rand(rng::AbstractRNG,
                   ::Random.SamplerType{Poly{D,T}}) where {D,T}
    n = rand(0:5)
    return Poly{D,T}(rand(rng, Term{D,T}, n))
end

@testset "Polynomials D=$D" for D in 0:Dmax
    T = Int

    for iter in 1:100
        n = zero(Poly{D,T})
        e = one(Poly{D,T})
        x = rand(Poly{D,T})
        y = rand(Poly{D,T})
        z = rand(Poly{D,T})
        a = T(rand(-10:10))
        b = T(rand(-10:10))
        i = rand(0:5)
        j = rand(0:5)

        # Ensure terms are nonzero
        for i in 1:length(x.terms)
            @test !iszero(x.terms[i])
        end

        # Ensure terms are sorted and unique
        for i in 2:length(x.terms)
            @test x.terms[i - 1].powers > x.terms[i].powers
            @test compare(x.terms[i - 1], x.terms[i]) < 0
        end

        @test n == n
        @test e == e
        @test e != n

        f(x) = 2 * x
        g(x) = x + 1
        h(x) = x^2
        @test map(identity, x) == x
        @test map(f ∘ identity, x) == map(f, x)
        @test map(identity ∘ f, x) == map(f, x)
        @test map(h ∘ g, map(f, x)) == map(h, map(g ∘ f, x))

        @test n + x == x
        @test x + n == x
        @test x + y == y + x
        @test x + (y + z) == (x + y) + z
        @test +x == x

        @test x + (-x) == n
        @test -(-x) == x
        @test -n == n
        @test x - y == x + (-y)

        @test zero(T) * x == n
        @test one(T) * x == x
        @test (-one(T)) * x == -x
        @test a * x == x * a
        @test (a * b) * x == a * (b * x)
        @test (a + b) * x == a * x + b * x
        @test a * (x + y) == a * x + a * y

        # Cannot test this with integers
        # if a != zero(T)
        #     @test (a * x) / a == x
        #     @test a \ (a * x) == x
        # end

        @test n * x == n
        @test x * n == n
        @test e * x == x
        @test x * e == x
        @test x * y == y * x
        @test x * (y * z) == (x * y) * z
        @test (a * x) * y == a * (x * y)
        @test x * (y * a) == (x * y) * a
        @test (x + y) * z == x * z + y * z

        @test x^0 == e
        @test x^1 == x
        @test x^(i + j) == x^i * x^j

        if D > 0
            dir = rand(1:D)

            dx = deriv(x, dir)
            @test deriv(n, dir) == n
            @test deriv(e, dir) == n
            @test deriv(unit(Poly{D,T}, dir), dir) == e
            @test deriv(x + y, dir) == deriv(x, dir) + deriv(y, dir)
            @test deriv(a * x, dir) == a * deriv(x, dir)

            @test deriv(x * y, dir) == deriv(x, dir) * y + x * deriv(y, dir)

            κx = koszul(x, dir)
            @test koszul(n, dir) == n
            @test koszul(e, dir) == unit(Poly{D,T}, dir)
            @test koszul(x + y, dir) == koszul(x, dir) + koszul(y, dir)
            @test koszul(a * x, dir) == a * koszul(x, dir)

            # This is not true
            # @test koszul(x * y, dir) == koszul(x, dir) * y + x * koszul(y, dir)
        end
    end
end

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

# Random polynomial spaces
function Base.rand(rng::AbstractRNG,
                   ::Random.SamplerType{PolySpace{D,T}}) where {D,T}
    n = rand(0:5)
    ps = rand(rng, Poly{D,T}, n)
    ps = filter(p -> !iszero(p), ps)
    return PolySpace{D,T}(ps)
end

@testset "Polynomial spaces D=$D" for D in 0:Dmax
    T = Int

    for iter in 1:100
        ps = rand(PolySpace{D,T})
        qs = rand(PolySpace{D,T})
        rs = rand(PolySpace{D,T})
        ns = PolySpace{D,T}()
        a = rand(-10:10)
        b = rand(-10:10)

        #TODO
        # # Ensure polynomials are nonzero
        # for i in 1:length(ps.polys)
        #     @test !iszero(ps.polys[i])
        # end

        #TODO
        # # Ensure polynomials are sorted and unique
        # for i in 2:length(ps.polys)
        #     @show i ps ps.polys[i - 1] ps.polys[i]
        #     @test compare(ps.polys[i - 1], ps.polys[i]) < 0
        # end

        @test ps == ps
        @test qs == qs
        @test (issubset(ps, qs) && issubset(qs, ps)) == (ps == qs)

        @test isempty(ns)
        @test isempty(ps) == (length(ps) == 0)

        @test union(ps, ns) == ps
        @test union(ps, qs) == union(qs, ps)
        @test union(union(ps, qs), rs) == union(ps, union(qs, rs))

        @test setdiff(ps, ns) == ps
        @test setdiff(setdiff(ps, qs), rs) == setdiff(ps, union(qs, rs))

        @test iszero(ns)
        # @test ps + ns == ps
        # @test ps + qs == qs + ps
        # @test (ps + qs) + rs == ps + (qs + rs)

        for dir in 1:D
            @test deriv(union(ps, qs), dir) ==
                  union(deriv(ps, dir), deriv(qs, dir))
            @test koszul(union(ps, qs), dir) ==
                  union(koszul(ps, dir), koszul(qs, dir))
        end
    end
end

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

@testset "Polynomial forms as vector space D=$D R=$R" for D in 0:Dmax, R in 0:D
    T = Int

    for iter in 1:100
        n = zero(Form{D,R,Poly{D,T}})
        x = rand(Form{D,R,Poly{D,T}})
        x′ = rand(Form{D,R,Poly{D,T}})
        y = rand(Form{D,R,Poly{D,T}})
        z = rand(Form{D,R,Poly{D,T}})
        a = T(rand(-10:10))
        b = T(rand(-10:10))
        pn = zero(Poly{D,T})
        pe = one(Poly{D,T})
        p = rand(Poly{D,T})
        q = rand(Poly{D,T})

        @test n == n

        @test n + x == x
        @test x + n == x
        @test x + y == y + x
        @test x + (y + z) == (x + y) + z
        @test +x == x

        @test x + (-x) == n
        @test -(-x) == x
        @test -n == n
        @test x - y == x + (-y)

        @test zero(T) * x == n
        @test one(T) * x == x
        @test (-one(T)) * x == -x
        @test a * x == x * a
        @test (a * b) * x == a * (b * x)
        @test (a + b) * x == a * x + b * x
        @test a * (x + y) == a * x + a * y

        @test pn * x == n
        @test pe * x == x
        @test (-pe) * x == -x
        @test p * x == x * p
        @test (p * q) * x == p * (q * x)
        @test (p + q) * x == p * x + q * x
        @test p * (x + y) == p * x + p * y
    end
end

@testset "Polynomial forms as ring D=$D" for D in 0:Dmax
    T = Int

    for iter in 1:100
        Rn = rand(0:D)
        Rx = rand(0:D)
        Ry = rand(0:D)
        Rz = rand(0:D)
        n = zero(Form{D,Rn,Poly{D,T}})
        e = one(Form{D,0,Poly{D,T}})
        x = rand(Form{D,Rx,Poly{D,T}})
        x′ = rand(Form{D,Rx,Poly{D,T}})
        y = rand(Form{D,Ry,Poly{D,T}})
        z = rand(Form{D,Rz,Poly{D,T}})
        a = T(rand(-10:10))
        p = rand(Poly{D,T})

        @test e == e
        @test iszero(n)
        @test isone(e)
        @test !iszero(e)
        @test !isone(n)

        if Rn + Rx <= D
            @test iszero(n ∧ x)
            @test iszero(x ∧ n)
        end
        @test e ∧ x == x
        @test x ∧ e == x
        if Rx + Ry <= D
            @test x ∧ y == bitsign(Rx * Ry) * (y ∧ x)
        end
        if Rx + Ry + Rz <= D
            @test x ∧ (y ∧ z) == (x ∧ y) ∧ z
        end
        if Rx + Ry <= D
            @test (a * x) ∧ y == a * (x ∧ y)
            @test x ∧ (y * a) == (x ∧ y) * a
            @test (p * x) ∧ y == p * (x ∧ y)
            @test x ∧ (y * p) == (x ∧ y) * p
            @test (x + x′) ∧ y == x ∧ y + x′ ∧ y
        end
    end
end

@testset "Derivatives of polynomial forms D=$D" for D in 0:Dmax
    T = Int

    for iter in 1:100
        Rx = rand(0:D)
        Ry = rand(0:D)
        x = rand(Form{D,Rx,Poly{D,T}})
        x2 = rand(Form{D,Rx,Poly{D,T}})
        y = rand(Form{D,Ry,Poly{D,T}})
        a = T(rand(-10:10))

        if Rx <= D - 1
            @test deriv(x + x2) == deriv(x) + deriv(x2)
            @test deriv(a * x) == a * deriv(x)
        end
        if Rx + Ry <= D - 1
            @test deriv(x ∧ y) == deriv(x) ∧ y + bitsign(Rx) * x ∧ deriv(y)
        end
        if Rx <= D - 2
            @test iszero(deriv(deriv(x)))
        end

        if 1 <= Rx
            @test koszul(x + x2) == koszul(x) + koszul(x2)
            @test koszul(a * x) == a * koszul(x)
        end
        if 1 <= min(Rx, Ry) && Rx + Ry <= D
            @test koszul(x ∧ y) == koszul(x) ∧ y + bitsign(Rx) * x ∧ koszul(y)
        end
        if 2 <= Rx
            @test iszero(koszul(koszul(x)))
        end

        if D >= 1
            Rz = rand(0:D)
            p = rand(0:5)
            function mkterm()
                dims = rand(1:D, p)
                powers = zeros(D)
                for d in dims
                    powers[d] += 1
                end
                return Term{D,T}(SVector{D,Int}(powers), T(rand(-10:10)))
            end
            function mkpoly()
                n = rand(0:5)
                return Poly{D,T}(Term{D,T}[mkterm() for i in 1:n])
            end
            function mkform()
                N = binomial(Val(D), Val(Rz))
                return Form{D,Rz}(ntuple(_ -> mkpoly(), N))
            end
            z = mkform()

            # Douglas Arnold, Richard Falk, Ragnar Winther, "Finite
            # element exterior calculus, homological techniques, and
            # applications", Acta Numerica 15, 1-155 (2006),
            # DOI:10.1017/S0962492906210018, eqn. (3.9)
            if 1 <= Rz <= D - 1
                # This holds only for homogenous polynomial forms,
                # i.e. forms where all polynomials have the same
                # degree `p`
                @test deriv(koszul(z)) + koszul(deriv(z)) == (Rz + p) * z
            end
        end
    end
end

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

@testset "Convert polynomial forms to vectors D=$D R=$R" for D in
                                                             0:min(5, Dmax),
R in 0:D

    T = Int

    for iter in 1:100
        x = rand(Form{D,R,Poly{D,T}})
        y = rand(Form{D,R,Poly{D,T}})
        n = zero(Form{D,R,Poly{D,T}})
        a = rand(-10:10)

        @test maxpower(n) == 0
        p = max(maxpower(x), maxpower(y))

        @test iszero(form2vec(n, p))
        @test form2vec(n + x, p) == form2vec(x, p)
        @test form2vec(x + y, p) == form2vec(x, p) + form2vec(y, p)
        @test form2vec(a * x, p) == a * form2vec(x, p)

        nrows = length(form2vec(n, p))
        @test forms2mat([x], p) == reshape(form2vec(x, p), nrows, 1)
    end

    for iter in 1:10
        ncols = rand(1:5)
        xs = [rand(Form{D,R,Poly{D,T}}) for j in 1:ncols]
        ys = [rand(Form{D,R,Poly{D,T}}) for j in 1:ncols]
        ns = [zero(Form{D,R,Poly{D,T}}) for j in 1:ncols]
        a = rand(-10:10)

        @test maxpower(ns) == 0
        p = max(maxpower(xs), maxpower(ys))

        @test iszero(forms2mat(ns, p))
        @test forms2mat(ns + xs, p) == forms2mat(xs, p)
        @test forms2mat(xs + ys, p) == forms2mat(xs, p) + forms2mat(ys, p)
        @test forms2mat(a * xs, p) == a * forms2mat(xs, p)
    end
end

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

@testset "Bases D=$D R=$R p=$p" for D in 0:Dmax, R in 0:D, p in 0:4
    T = Int

    # feec-icerm-lecture3, page 23
    # n = D
    # r = p
    # k = R

    fb = full_basis(Basis{D,R,T}, p)
    @test length(fb.forms) == binomial(p + D, p) * binomial(D, R)

    hb = homogeneous_basis(Basis{D,R,T}, p)
    @test length(hb.forms) == binomial(p + D - 1, p) * binomial(D, R)
end

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

@testset "Complexes D=$D p=$p" for D in 0:Dmax, p in 1:max(2, 5 - D)
    T = Int

    pc = polynomial_complex(Val(D), Int, p)
    tpc = trimmed_polynomial_complex(Val(D), Int, p)
    epc = extended_trimmed_polynomial_complex(Val(D), Int, p)

    @test Set(keys(pc)) == Set(0:D)
    @test Set(keys(tpc)) == Set(0:D)
    @test Set(keys(epc)) == Set(0:D)

    for R in 0:D
        @test pc[R] isa Basis{D,R,T}
        @test tpc[R] isa Basis{D,R,T}
        @test epc[R] isa Basis{D,R,T}

        fb = full_basis(Basis{D,R,T}, p)
        fb1 = full_basis(Basis{D,R,T}, p - 1)

        if R != 0 && R != D
            @test issubset(fb1, tpc[R])
            @test fb1 != tpc[R]
        end
        if R == 0
            @test tpc[R] == pc[R]
        elseif R == D
            @test tpc[R] == fb1
        else
            @test issubset(tpc[R], fb)
            @test tpc[R] != fb
        end
        @test issubset(tpc[R], epc[R])
        if R == 0
            @test epc[R] == pc[R]
        elseif R == D
            # This does not seem to hold
            # @test epc[R] == fb1
        else
            @test issubset(epc[R], fb)
            @test epc[R] != pc[R]
        end

        if R < D
            @test issubset(deriv(pc[R]), pc[R + 1])
            @test issubset(deriv(tpc[R]), tpc[R + 1])
            # This does not seem to hold
            # This also means that epc is not actually complex!
            # @test issubset(deriv(epc[R]), epc[R + 1])
        end
        if R > 0
            @test issubset(koszul(pc[R]), pc[R - 1])
            @test issubset(koszul(tpc[R]), tpc[R - 1])
            # This does not seem to hold
            # @test issubset(koszul(epc[R]), epc[R - 1])
        end
    end
end

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

@testset "Whitney forms D=$D R=$R" for D in 0:Dmax, R in 0:D
    T = Int

    ϕ = whitney(Basis{D + 1,R,Int})
    @test length(ϕ.forms) == binomial(D + 1, R + 1)

    for n in 1:length(ϕ.forms)
        ϕₙ = ϕ.forms[n]
        ϕ′ = Basis{D + 1,R,T}([ϕ.forms[1:(n - 1)]; ϕ.forms[(n + 1):end]])
        @test !(ϕₙ in ϕ′)
    end
end

@testset "Convert from barycentric to Cartesian coordinates D=$D R=$R" for D in
                                                                           0:Dmax,
R in 0:D

    T = Int

    ϕ = whitney(Basis{D + 1,R,Int})
    @test length(ϕ.forms) == binomial(D + 1, R + 1)
    ϕx = barycentric2cartesian(ϕ)

    for n in 1:length(ϕx.forms)
        ϕₙ = ϕx.forms[n]
        ϕx′ = Basis{D,R,T}([ϕx.forms[1:(n - 1)]; ϕx.forms[(n + 1):end]])
        @test !(ϕₙ in ϕx′)
    end
end

@testset "Compare Whitney basis to trimmed p=1 basis D=$D R=$R" for D in 0:Dmax,
R in 0:D

    T = Int
    p = 1                       # Whitney basis is equivalent to p=1

    ϕλ = whitney(Basis{D + 1,R,Int})
    ϕ = barycentric2cartesian(ϕλ)

    tpc = trimmed_polynomial_complex(Val(D), Int, p)
    tpcR = tpc[R]

    @test ϕ == tpcR
end


module day03
using CircularArrays

include("inputs.jl")
ex = inputs.ex_d3()
input = inputs.input_d3()

# Part 1
function parsemap(input::String)
    map = split(input) .|> # array of strings "..##......."
    collect .|> # array of array of chars ['.', '.', '#', '#', ...]
    x->(x .=='#') |> # array of BitArray [0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0]
    x->hcat(x...)    # array of Bools 0  0  1  1  0  0  0  0  0  0  0
    map = vcat(map...) |>
    CircularArray
end

function count_trees(input, slope=(3,1))
    map = parsemap(input)
    count = 0
    x, y = 1, 1
    while y<=size(map)[1]
        count += map[y, x] ? 1 : 0
        x += slope[1]
        y += slope[2]
    end
    return count
end

@assert count_trees(ex) == 7
@info "Day 3, Part 1 answer: $(count_trees(input))"

# Part 2
function assess_slopes(input)
    slopes = [(1,1), (3,1), (5,1), (7,1), (1,2)]
    counts = [count_trees(input, slope) for slope in slopes]
    return reduce(*, counts)
end

@assert assess_slopes(ex) == 336
@info "Day 3, Part 2 answer: $(assess_slopes(input))"

end


"""
This can contain a maximum of `length(par)` parenting-relations
par is an array of `Int`, which is the index of the parent node.
"""
mutable struct DisjointSet
    par::Vector{Int}
    function DisjointSet(size)
        x = [i for i in 1:size]
        return new(x)
    end
end

"""
Find the ancestor of node `x`.
"""
function find(set::DisjointSet, x::Int)::Int
    if set.par[x] == x
        return x
    else
        return set.par[x] = find(set, set.par[x])
    end
end

function Base.merge!(set::DisjointSet, x::Int, y::Int)
    x = find(set, x)
    y = find(set, y)
    return set.par[x] = y
end


"""
    transition_matrices(m::Union{MDP,POMDP})
    transition_matrices(m; sparse=true)

Construct transition matrices for (PO)MDP m.

The returned object is an associative object (usually a Dict), where the keys are actions. Each value in this object is an AbstractMatrix where the row corresponds to the state index of s and the column corresponds to the state index of s'. The entry in the matrix is the probability of transitioning from state s to state s'.
"""
function transition_matrices(m::Union{MDP,POMDP}; sparse::Bool=true)
    transmats = POMDPModelTools.transition_matrix_a_s_sp(m)
    if !sparse
        transmats = [convert(Matrix, t) for t in transmats]
    end
    mtype = typeof(first(transmats))
    oa = ordered_actions(m)
    return Dict{actiontype(m), mtype}(oa[ai]=>transmats[ai] for ai in 1:length(actions(m)))
end

"""
    reward_vectors(m::Union{MDP, POMDP})

Construct reward vectors for (PO)MDP m.

The returned object is an associative object (usually a Dict), where the keys are actions. Each value in this object is an AbstractVector where the index corresponds to the state index of s and the entry is the reward for that state.
"""
function reward_vectors(m::Union{MDP,POMDP})
    d = Dict{actiontype(m), Vector{Float64}}()
    r = StateActionReward(m)
    for a in actions(m)
        rv = zeros(length(states(m)))
        for s in states(m)
            rv[stateindex(m, s)] = r(s, a)
        end
        d[a] = rv
    end
    return d
end


@testset "BernsteinPolynomial" begin
    @testset "evaluation" begin
        b = BernsteinPolynomial(1, 2, 3, 4)
        @test b(0) === 1
        @test b(1) === 4
    end

    @testset "utility" begin
        b = BernsteinPolynomial(1, 2, 3, 4)
        @test zero(b) === zero(typeof(b)) === BernsteinPolynomial(0, 0, 0, 0)
        @test constant(b) == 1
    end

    @testset "derivative" begin
        b = BernsteinPolynomial(1, 2, 3, 4)
        b′ = derivative(b)
        for t in range(0., 1., length=10)
            @test ForwardDiff.derivative(b, t) ≈ b′(t) atol=1e-10
        end
    end

    @testset "exponential integral" begin
        b = BernsteinPolynomial(1, 2, 3, 4)
        c = 5
        for t in range(0., 1., length=10)
            @test ForwardDiff.derivative(t -> exponential_integral(b, c, t), t) ≈ b(t) * exp(c * t) atol=1e-10
        end
        @test exponential_integral(b, c) ≈ exponential_integral(b, c, 1.0) atol=1e-10
    end

    @testset "arithmetic" begin
        b = BernsteinPolynomial(1, 2, 3, 4)
        c = 5
        for t in range(0., 1., length=10)
            @test (b + c)(t) ≈ b(t) + c
            @test (b - c)(t) ≈ b(t) - c
            @test (c + b)(t) ≈ c + b(t)
            @test (c - b)(t) ≈ c - b(t)
        end
        b1 = BernsteinPolynomial(1, 2, 3, 4)
        b2 = BernsteinPolynomial(2, 3, 4, 5)
        for t in range(0., 1., length=10)
            @test (b1 + b2)(t) ≈ b1(t) + b2(t)
            @test (b1 - b2)(t) ≈ b1(t) - b2(t)
        end
    end

    @testset "scaling" begin
        b1 = BernsteinPolynomial(2, 3, 4)
        @test b1 * 4 === 4 * b1 === BernsteinPolynomial(8, 12, 16)
        @test b1 / 2 === BernsteinPolynomial(2 / 2, 3 / 2, 4 / 2)
        for t in range(0., 1., length=10)
            @test (b1 * 4)(t) ≈ b1(t) * 4
            @test (b1 / 2)(t) ≈ b1(t) / 2
        end
        b2 = BernsteinPolynomial(ntuple(i -> 21 - i, Val(20)))
        allocs = let b2 = b2
            @allocated b2 * 4
        end
        @test allocs == 0
    end

    @testset "BernsteinPolynomial to Polynomial" begin
        b = BernsteinPolynomial(5, 4, 3, 2, 1, 10)
        p = Polynomial(b)
        for t in range(0., 1., length=10)
            @test b(t) ≈ p(t)
        end
        allocs = let b = b
            @allocated Polynomial(b)
        end
        @test allocs == 0
    end
end


# ---
# title: 1012. Numbers With Repeated Digits
# id: problem1012
# author: Tian Jun
# date: 2020-10-31
# difficulty: Hard
# categories: Math, Dynamic Programming
# link: <https://leetcode.com/problems/numbers-with-repeated-digits/description/>
# hidden: true
# ---
#
# Given a positive integer `N`, return the number of positive integers less than
# or equal to `N` that have at least 1 repeated digit.
#
#
#
# **Example 1:**
#
#
#
#     Input: 20
#     Output: 1
#     Explanation: The only positive number (<= 20) with at least 1 repeated digit is 11.
#
#
# **Example 2:**
#
#
#
#     Input: 100
#     Output: 10
#     Explanation: The positive numbers (<= 100) with atleast 1 repeated digit are 11, 22, 33, 44, 55, 66, 77, 88, 99, and 100.
#
#
# **Example 3:**
#
#
#
#     Input: 1000
#     Output: 262
#
#
#
#
# **Note:**
#
#   1. `1 <= N <= 10^9`
#
#
## @lc code=start
using LeetCode

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


"""
    MPS

A finite size matrix product state type.
Keeps track of the orthogonality center.
"""
mutable struct MPS <: AbstractMPS
  length::Int
  data::Vector{ITensor}
  llim::Int
  rlim::Int
  function MPS(N::Int,
               A::Vector{<:ITensor},
               llim::Int = 0,
               rlim::Int = N+1)
    new(N, A, llim, rlim)
  end
end

"""
    MPS()

Construct an empty MPS with zero sites.
"""
MPS() = MPS(0, Vector{ITensor}(), 0, 0)

"""
    MPS(N::Int)

Construct an MPS with N sites with default constructed
ITensors.
"""
MPS(N::Int) = MPS(N, Vector{ITensor}(undef, N))

"""
    MPS(::Type{T<:Number}, sites)

Construct an MPS from a collection of indices
with element type T.
"""
function MPS(::Type{T}, sites) where {T<:Number}
  N = length(sites)
  v = Vector{ITensor}(undef, N)
  l = [Index(1, "Link,l=$ii") for ii=1:N-1]
  for ii in eachindex(sites)
    s = sites[ii]
    if ii == 1
      v[ii] = ITensor(T,l[ii], s)
    elseif ii == N
      v[ii] = ITensor(T,l[ii-1], s)
    else
      v[ii] = ITensor(T,l[ii-1],s,l[ii])
    end
  end
  return MPS(N, v)
end

"""
    MPS(sites)

Construct an MPS from a collection of indices
with element type Float64.
"""
MPS(sites) = MPS(Float64, sites)

function randomU(s1,s2)
  if !hasqns(s1) && !hasqns(s2)
    mdim = dim(s1)*dim(s2)
    RM = randn(mdim,mdim)
    Q,_ = NDTensors.qr_positive(RM)
    G = itensor(Q,dag(s1),dag(s2),s1',s2')
  else
    M = randomITensor(QN(),s1',s2',dag(s1),dag(s2))
    U,S,V = svd(M,(s1',s2'))
    u = commonind(U,S)
    v = commonind(S,V)
    G = (U*delta(dag(u),v))*V
  end
  return G
end

function randomizeMPS!(M::MPS, sites, linkdim=1)
  N = length(sites)
  c = div(N,2)
  max_pass = 100
  for pass=1:max_pass,half=1:2
    if half==1
      (db,brange) = (+1, 1:1:N-1)
    else
      (db,brange) = (-1, N:-1:2)
    end
    for b=brange
      s1 = sites[b]
      s2 = sites[b+db]
      G = randomU(s1,s2)
      T = noprime(G*M[b]*M[b+db])
      rinds = uniqueinds(M[b],M[b+db])
      U,S,V = svd(T,rinds;maxdim=linkdim)
      M[b] = U
      M[b+db] = S*V
      M[b+db] /= norm(M[b+db])
    end
    if half==2 && dim(commonind(M[c],M[c+1])) >= linkdim
      break
    end
  end
  setleftlim!(M, 0)
  setrightlim!(M, 2)
  if dim(commonind(M[c],M[c+1])) < linkdim
    error("MPS center bond dim less than requested")
  end
end

function randomCircuitMPS(sites,linkdim::Int;kwargs...)::MPS
  N = length(sites)
  M = MPS(N)

  l = Vector{Index}(undef,N)

  d = dim(sites[N])
  chi = min(linkdim,d)
  l[N-1] = Index(chi,"Link,l=$(N-1)")
  O = NDTensors.random_orthog(chi,d)
  M[N] = itensor(O,l[N-1],sites[N])

  for j=N-1:-1:2
    chi *= dim(sites[j])
    chi = min(linkdim,chi)
    l[j-1] = Index(chi,"Link,l=$(j-1)")
    O = NDTensors.random_orthog(chi,dim(sites[j])*dim(l[j]))
    T = reshape(O,(chi,dim(sites[j]),dim(l[j])))
    M[j] = itensor(T,l[j-1],sites[j],l[j])
  end

  O = NDTensors.random_orthog(1,dim(sites[1])*dim(l[1]))
  l0 = Index(1,"Link,l=0")
  T = reshape(O,(1,dim(sites[1]),dim(l[1])))
  M[1] = itensor(T,l0,sites[1],l[1])
  M[1] *= setelt(l0=>1)

  M.llim = 0
  M.rlim = 2

  return M
end

"""
    randomMPS(::Type{T<:Number}, sites; linkdim=1)

Construct a random MPS with link dimension `linkdim` of
type `T`.
"""
function randomMPS(::Type{T}, sites, linkdim::Int=1) where {T<:Number}
  if hasqns(sites[1])
    error("initial state required to use randomMPS with QNs")
  end

  # For non-QN-conserving MPS, instantiate
  # the random MPS directly as a circuit:
  return randomCircuitMPS(sites,linkdim)
end

"""
    randomMPS(sites; linkdim=1)

Construct a random MPS with link dimension `linkdim` of
type `Float64`.
"""
randomMPS(sites, linkdim::Int=1) = randomMPS(Float64, sites, linkdim)

"""
    randomMPS(sites,state; linkdim=1)

Construct a real, random MPS with link dimension `linkdim`,
made by randomizing an initial product state specified by
`state`.
"""
function randomMPS(sites,state,linkdim::Int=1)::MPS
  M = productMPS(sites,state)
  if linkdim > 1
    randomizeMPS!(M,sites,linkdim)
  end
  return M
end

"""
    productMPS(::Type{T<:Number}, ivals::Vector{<:IndexVal})

Construct a product state MPS with element type `T` and
nonzero values determined from the input IndexVals.
"""
function productMPS(::Type{T},
                    ivals::Vector{<:IndexVal}) where {T<:Number}
  N = length(ivals)
  M = MPS(N)
  if hasqns(ind(ivals[1]))
    links = [Index(QN()=>1;tags="Link,l=$n") for n=1:N]
  else
    links = [Index(1,"Link,l=$n") for n=1:N]
  end
  M[1] = ITensor(T,ind(ivals[1]),links[1])
  M[1][ivals[1],links[1](1)] = 1.0
  for n=2:N-1
    s = ind(ivals[n])
    M[n] = ITensor(T,dag(links[n-1]),s,links[n])
    M[n][links[n-1](1),ivals[n],links[n](1)] = 1.0
  end
  M[N] = ITensor(T,dag(links[N-1]),ind(ivals[N]))
  M[N][links[N-1](1),ivals[N]] = 1.0
  return M
end

"""
    productMPS(ivals::Vector{<:IndexVal})

Construct a product state MPS with element type `Float64` and
nonzero values determined from the input IndexVals.
"""
productMPS(ivals::Vector{<:IndexVal}) = productMPS(Float64,
                                                   ivals::Vector{<:IndexVal})

function productMPS(::Type{T},
                    sites,
                    states) where {T<:Number}
  if length(sites) != length(states)
    throw(DimensionMismatch("Number of sites and and initial states don't match"))
  end
  ivals = [state(sites[n],states[n]) for n=1:length(sites)]
  return productMPS(T, ivals)
end

productMPS(sites, states) = productMPS(Float64, sites, states)

function linkind(M::MPS, j::Int)
  N = length(M)
  j ≥ length(M) && error("No link index to the right of site $j (length of MPS is $N)")
  li = commonind(M[j],M[j+1])
  if isnothing(li)
    error("linkind: no MPS link index at link $j")
  end
  return li
end

function siteind(M::MPS, j::Int)
  N = length(M)
  if j == 1
    si = uniqueind(M[j],M[j+1])
  elseif j == N
    si = uniqueind(M[j],M[j-1])
  else
    si = uniqueind(M[j],M[j-1],M[j+1])
  end
  return si
end

function siteinds(M::MPS)
  return [siteind(M,j) for j in 1:length(M)]
end

function replacesiteinds!(M::MPS, sites)
  for j in eachindex(M)
    sj = siteind(M,j)
    replaceind!(M[j], sj, sites[j])
  end
  return M
end

"""
    dot(psi::MPS, phi::MPS; make_inds_match = true)
    inner(psi::MPS, phi::MPS; make_inds_match = true)

Compute <psi|phi>.

If `make_inds_match = true`, the function attempts to make
the site indices match before contracting (so for example, the
inputs can have different site indices, as long as they
have the same dimensions or QN blocks).
"""
function LinearAlgebra.dot(M1::MPS, M2::MPS; make_inds_match::Bool = true)::Number
  N = length(M1)
  if length(M2) != N
    throw(DimensionMismatch("inner: mismatched lengths $N and $(length(M2))"))
  end
  M1dag = dag(M1)
  if make_inds_match
    simlinkinds!(M1dag)
    replacesiteinds!(M1dag, siteinds(M2))
  end
  O = M1dag[1]*M2[1]
  for j in eachindex(M1)[2:end]
    O = (O*M1dag[j])*M2[j]
  end
  return O[]
end

inner(M1::MPS, M2::MPS; kwargs...) = dot(M1, M2; kwargs...)

"""
    replacebond!(M::MPS, b::Int, phi::ITensor; kwargs...)

Factorize the ITensor `phi` and replace the ITensors
`b` and `b+1` of MPS `M` with the factors. Choose
the orthogonality with `ortho="left"/"right"`.
"""
function replacebond!(M::MPS,
                      b::Int,
                      phi::ITensor;
                      kwargs...)
  ortho::String = get(kwargs, :ortho, "left")
  which_decomp::Union{String, Nothing} = get(kwargs, :which_decomp, nothing)
  normalize::Bool = get(kwargs, :normalize, false)

  # Deprecated keywords
  if haskey(kwargs, :dir)
    error("""dir keyword in replacebond! has been replaced by ortho.
          Note that the options are now the same as factorize, so use `left` instead of `fromleft` and `right` instead of `fromright`.""")
  end

  L,R,spec = factorize(phi,inds(M[b]); which_decomp = which_decomp,
                                       tags = tags(linkind(M,b)),
                                       kwargs...)
  M[b]   = L
  M[b+1] = R
  if ortho == "left"
    leftlim(M) == b-1 && setleftlim!(M, leftlim(M)+1)
    rightlim(M) == b+1 && setrightlim!(M, rightlim(M)+1)
    normalize && (M[b+1] ./= norm(M[b+1]))
  elseif ortho == "right"
    leftlim(M) == b && setleftlim!(M, leftlim(M)-1)
    rightlim(M) == b+2 && setrightlim!(M, rightlim(M)-1)
    normalize && (M[b] ./= norm(M[b]))
  else
    error("In replacebond!, got ortho = $ortho, only currently supports `left` and `right`.")
  end
  return spec
end

"""
    sample!(m::MPS)

Given a normalized MPS m, returns a `Vector{Int}`
of `length(m)` corresponding to one sample
of the probability distribution defined by
squaring the components of the tensor
that the MPS represents. If the MPS does
not have an orthogonality center,
orthogonalize!(m,1) will be called before
computing the sample.
"""
function sample!(m::MPS)
  orthogonalize!(m, 1)
  return sample(m)
end

"""
    sample(m::MPS)

Given a normalized MPS m with `orthocenter(m)==1`,
returns a `Vector{Int}` of `length(m)`
corresponding to one sample of the
probability distribution defined by
squaring the components of the tensor
that the MPS represents
"""
function sample(m::MPS)
  N = length(m)

  if orthocenter(m) != 1
    error("sample: MPS m must have orthocenter(m)==1")
  end
  if abs(1.0-norm(m[1])) > 1E-8
    error("sample: MPS is not normalized, norm=$(norm(m[1]))")
  end

  result = zeros(Int,N)
  A = m[1]

  for j=1:N
    s = siteind(m,j)
    d = dim(s)
    # Compute the probability of each state
    # one-by-one and stop when the random
    # number r is below the total prob so far
    pdisc = 0.0
    r = rand()
    # Will need n,An, and pn below
    n = 1
    An = ITensor()
    pn = 0.0
    while n <= d
      projn = ITensor(s)
      projn[s[n]] = 1.0
      An = A*projn
      pn = real(scalar(dag(An)*An))
      pdisc += pn
      (r < pdisc) && break
      n += 1
    end

    result[j] = n

    if j < N
      A = m[j+1]*An
      A *= (1.0/sqrt(pn))
    end
  end
  return result
end


using Test

using ChipSort
using SIMD

include("testutils.jl")

@testset "Transpose" begin

# types_to_test = [UInt8, UInt16, UInt32, UInt64, Float32, Float64]
types_to_test = [UInt16, Float64]

function test_transpose(T, V, J, K)
    data = randa(T, V*J*K)
    sol = permutedims(reshape(data, V, K, J), [1,3,2])[:]
    # seed=find_transpose_cycles(Val(J), Val(K))
    # transpose!(aa, Val(V), Val(J), Val(K), seed)
    transpose!(data, Val(V), Val(J), Val(K))
    @test sol == data
end

@testset for T in types_to_test
    for V in 2 .^ (1:2:5)
        for J in 2 .^ (1:2:5)
            for K in 2 .^ (2:4:8)
                test_transpose(T, V, J, K)
            end
        end
    end
end

function test_transpose_vecs(T, N, L)
    uu = ntuple(l->Vec(tuple([convert(T, mod(l*N+n-N, 2^(8*T.size-2))) for n in 1:N]...)), L)
    ttref = ntuple(l->Vec(tuple([convert(T, mod(n*N+l-N, 2^(8*T.size-2))) for n in 1:L]...)), N)
    tthat = transpose_vecs(uu...)
    @test all(all(j) for j in (tthat.==ttref))
end

@testset for T in types_to_test
    for N in 2 .^ (1:2:5)
        for L in 2 .^ (1:2:5)
            test_transpose_vecs(T, N, L)
        end
    end
end

end


using FixedPointSinCosApproximations
using Test

include("second_order.jl")
include("third_order.jl")
include("fourth_order.jl")
include("fifth_order.jl")
include("sixth_order.jl")
include("approximations.jl")


"""
    Manifold
A manifold type. The `Manifold` is used to dispatch to different exponential
and logarithmic maps as well as other function on manifold.
"""
abstract type Manifold end

abstract type AbstractRetractionMethod end

"""
    ExponentialRetraction
Retraction using the exponential map.
"""
struct ExponentialRetraction <: AbstractRetractionMethod end

"""
    retract!(M::Manifold, y, x, v, [t=1], [method::AbstractRetractionMethod=ExponentialRetraction()])
Retraction (cheaper, approximate version of exponential map) of tangent
vector `t*v` at point `x` from manifold `M`.
Result is saved to `y`.
Retraction method can be specified by the last argument. Please look at the
documentation of respective manifolds for available methods.
"""
retract!(M::Manifold, y, x, v, method::ExponentialRetraction) = exp!(M, y, x, v)

retract!(M::Manifold, y, x, v) = retract!(M, y, x, v, ExponentialRetraction())

retract!(M::Manifold, y, x, v, t::Real) = retract!(M, y, x, t*v)

retract!(M::Manifold, y, x, v, t::Real, method::AbstractRetractionMethod) = retract!(M, y, x, t*v, method)
struct Euclidean{T<:Tuple} <: Manifold where {T} end

Euclidean(n::Int) = Euclidean{Tuple{n}}()
Euclidean(m::Int, n::Int) = Euclidean{Tuple{m,n}}()

function representation_size(::Euclidean{Tuple{n}}) where {n}
    return (n,)
end

function representation_size(::Euclidean{Tuple{m,n}}) where {m,n}
    return (m,n)
end

@generated manifold_dimension(::Euclidean{T}) where {T} = *(T.parameters...)

exp!(M::Euclidean, y, x, v) = (y .= x .+ v)

log!(M::Euclidean, v, x, y) = (v .= y .- x)

function zero_tangent_vector!(M::Euclidean, v, x)
    fill!(v, 0)
    return v
end

project_point!(M::Euclidean, x) = x

function project_tangent!(M::Euclidean, w, x, v)
    w .= v
    return w
end
retract(M::Euclidean, x, v, t::Real) = x+t*v

function retract(M::Manifold, x, v, method::AbstractRetractionMethod)
    xr = similar_result(M, retract, x, v)
    retract!(M, xr, x, v, method)
    return xr
end

function retract(M::Manifold, x, v)
    xr = similar_result(M, retract, x, v)
    retract!(M, xr, x, v)
    return xr
end

retract(M::Manifold, x, v, t::Real) = retract(M, x, t*v)

retract(M::Manifold, x, v, t::Real, method::AbstractRetractionMethod) = retract(M, x, t*v, method)


# """
#     strain2x2tto3v!(v::FVec{T}, t::FMat{T}) where {T}

# Convert a matrix of 2x2 strain components  into a 3-component vector.
# """
# function strain2x2tto3v!(v::FVec{T}, t::FMat{T}) where {T}
#     v[1] = t[1,1];
#     v[2] = t[2,2];
#     v[3] = t[1,2] + t[2,1];
#     return v
# end

# """
#     strain3vto2x2t!(t::FMat{T}, v::FVec{T}) where {T}

# Convert a strain 3-vector to a  matrix of 2x2 strain components (symmetric tensor).
# """
# function strain3vto2x2t!(t::FMat{T}, v::FVec{T}) where {T}
#     t[1,1] = v[1];
#     t[2,2] = v[2];
#     t[1,2] = v[3]/2.;
#     t[2,1] = v[3]/2.;
#     return t
# end

# """
#     strain3x3tto6v!(v::FVec{T}, t::FMat{T}) where {T}

# Convert a matrix of 3x3 strain components to a 6-component vector.
# """
# function strain3x3tto6v!(v::FVec{T}, t::FMat{T}) where {T}
#     v[1] = t[1,1];
#     v[2] = t[2,2];
#     v[3] = t[3,3];
#     v[4] = t[1,2] + t[2,1];
#     v[5] = t[1,3] + t[3,1];
#     v[6] = t[3,2] + t[2,3];
#     return v
# end

# """
#     strain6vto3x3t!(t::FMat{T}, v::FVec{T}) where {T}

# Convert a strain 6-vector to a  matrix of 3x3 strain components (symmetric tensor)..
# """
# function strain6vto3x3t!(t::FMat{T}, v::FVec{T}) where {T}
#     t[1,1] = v[1];
#     t[2,2] = v[2];
#     t[3,3] = v[3];
#     t[1,2] = v[4]/2.;
#     t[2,1] = v[4]/2.;
#     t[1,3] = v[5]/2.;
#     t[3,1] = v[5]/2.;
#     t[3,2] = v[6]/2.;
#     t[2,3] = v[6]/2.;
#     return t
# end

# """
#     strain3x3tto9v!(v::FVec{T}, t::FMat{T}) where {T}

# Convert a matrix of 3x3 strain components to a 9-component vector.

# The  strain components are in the order
#      ex, ey, ez, gxy, gyx, gyz, gzy, gxz, gzx
# """
# function strain3x3tto9v!(v::FVec{T}, t::FMat{T}) where {T}
#     v[1] = t[1,1];
#     v[2] = t[2,2];
#     v[3] = t[3,3];
#     v[4] = t[1,2];
#     v[5] = t[2,1];
#     v[6] = t[2,3];
#     v[7] = t[3,2];
#     v[8] = t[1,3];
#     v[9] = t[3,1];
#     return v
# end

# """
#     strain9vto3x3t!(v::FVec{T}, t::FMat{T}) where {T}

# Convert a matrix of 3x3 strain components to a 9-component vector.

# The  strain components are in the order
#      ex, ey, ez, gxy, gyx, gyz, gzy, gxz, gzx
# """
# function strain9vto3x3t!(t::FMat{T}, v::FVec{T}) where {T}
#     t[1,1] = v[1];
#     t[2,2] = v[2];
#     t[3,3] = v[3];
#     t[1,2] = v[4];
#     t[2,1] = v[5];
#     t[2,3] = v[6];
#     t[3,2] = v[7];
#     t[1,3] = v[8];
#     t[3,1] = v[9];
#     return t
# end

# """
#     stress2x2to3v!(v::FVec{T}, t::FMat{T})

# Convert a symmetric matrix of 2x2 stress components to a 3-component vector.
# """
# function stress2x2to3v!(v::FVec{T}, t::FMat{T}) where {T}
#     v[1] = t[1,1];
#     v[2] = t[2,2];
#     v[3] = 0.5*(t[1,2] + t[2,1]);
#     return v
# end

# """
#     stress3vto2x2t!(t::FMat{T}, v::FVec{T})

# Convert a 3-vector to a  matrix of 2x2 stress components (symmetric tensor).
# """
# function stress3vto2x2t!(t::FMat{T}, v::FVec{T}) where {T}
#     t[1,1] = v[1];
#     t[2,2] = v[2];
#     t[1,2] = v[3];
#     t[2,1] = v[3];
#     return t
# end

# """
#     stress3vto3x3t!(t::FMat{T}, v::FVec{T}) where {T}

# Convert a 3-vector to a matrix of 3x3 stress components (symmetric tensor).
# """
# function stress3vto3x3t!(t::FMat{T}, v::FVec{T}) where {T}
#     t[1,1] = v[1];
#     t[2,2] = v[2];
#     t[1,2] = v[3];
#     t[2,1] = v[3];
#     return t
# end

# """
#     stress4vto3x3t!(t::FMat{T}, v::FVec{T}) where {T}

# Convert a 4-vector to a  matrix of 3x3 stress components (tensor).

# Convert a 4-vector to a *symmetric*
# matrix of 3x3 stress components (tensor).  This is
# conversion routine that would be useful for plane strain or
# axially symmetric conditions.
# The stress vector components need to be ordered as:
#     sigmax, sigmay, tauxy, sigmaz,
# which is the ordering used for the plane-strain model reduction.
# Therefore, for axially symmetric analysis the components need to be
# reordered, as from the constitutive equation they come out
# as sigmax, sigmay, sigmaz, tauxy.
# """
# function stress4vto3x3t!(t::FMat{T}, v::FVec{T}) where {T}
#     t[1,1] = v[1];
#     t[2,2] = v[2];
#     t[1,2] = v[3];
#     t[2,1] = v[3];
#     t[3,3] = v[4];
#     return t
# end

# """
#     stress6vto3x3t!(t::FMat{T}, v::FVec{T}) where {T}

# Convert a 6-vector to a  matrix of 3x3 stress components (symmetric tensor).
# """
# function stress6vto3x3t!(t::FMat{T}, v::FVec{T}) where {T}
#     t[1,1] = v[1];
#     t[2,2] = v[2];
#     t[3,3] = v[3];
#     t[1,2] = v[4];
#     t[2,1] = v[4];
#     t[1,3] = v[5];
#     t[3,1] = v[5];
#     t[3,2] = v[6];
#     t[2,3] = v[6];
#     return t
# end

# """
#     stress3x3tto6v!(v::FVec{T}, t::FMat{T}) where {T}

# Convert a matrix of 3x3 stress components to a 6-component vector.
# """
# function stress3x3tto6v!(v::FVec{T}, t::FMat{T}) where {T}
#     v[1] = t[1,1];
#     v[2] = t[2,2];
#     v[3] = t[3,3];
#     v[4] = 1/2.0*(t[1,2] + t[2,1]);
#     v[5] = 1/2.0*(t[1,3] + t[3,1]);
#     v[6] = 1/2.0*(t[3,2] + t[2,3]);
#     return v
# end

# """
#     strain9vto6v!(t::FVec{T}, v::FVec{T}) where {T}

# Convert a strain 9-vector to a  strain 6-vector components (tensor).
# """
# function strain9vto6v!(t::FVec{T}, v::FVec{T}) where {T}
#     t[1] = v[1];
#     t[2] = v[2];
#     t[3] = v[3];
#     t[4] = v[4]+v[5];
#     t[5] = v[8]+v[9];
#     t[6] = v[6]+v[7];
#     return t
# end

# """
#     strain6vto9v!(t::FVec{T}, v::FVec{T}) where {T}

# Convert a strain 6-vector to a  strain 9-vector components (tensor).

# The  strain components are in the order
#      ex, ey, ez, gxy/2, gxy/2, gyz/2, gyz/2, gxz/2, gxz/2
# """
# function strain6vto9v!(t::FVec{T}, v::FVec{T}) where {T}
#     t[1] = v[1];
#     t[2] = v[2];
#     t[3] = v[3];
#     t[4] = v[4]/2.;
#     t[5] = v[4]/2.;
#     t[6] = v[6]/2.;
#     t[7] = v[6]/2.;
#     t[8] = v[5]/2.;
#     t[9] = v[5]/2.;
#     return t
# end

# """
#     stress9vto6v!(t::FVec{T}, v::FVec{T}) where {T}

# Convert a stress 9-vector (tensor) to a  stress 6-vector components.

# The  stress components are in the order
#      sigx, sigy, sigz, tauxy, tauxy, tauyz, tauyz, tauxz, tauxz
# """
# function stress9vto6v!(t::FVec{T}, v::FVec{T}) where {T}
#     t[1] = v[1];
#     t[2] = v[2];
#     t[3] = v[3];
#     t[4] = v[4];
#     t[5] = v[8];
#     t[6] = v[6];
#     return t
# end

# """
#     stress6vto9v!(t::FVec{T}, v::FVec{T}) where {T}

# Convert a stress 6-vector to a  stress 9-vector components (tensor).

# The  stress components are in the order
#      sigx, sigy, sigz, tauxy, tauxy, tauyz, tauyz, tauxz, tauxz
# """
# function stress6vto9v!(t::FVec{T}, v::FVec{T}) where {T}
#     t[1] = v[1];
#     t[2] = v[2];
#     t[3] = v[3];
#     t[4] = v[4];
#     t[5] = v[4];
#     t[6] = v[6];
#     t[7] = v[6];
#     t[8] = v[5];
#     t[9] = v[5];
#     return t
# end

# function stressvectorrotation{MR<:DeforModelRed2DStress}(::Type{MR},
#   Rm::FMat{T})
#
#     # Rm = columns are components of 'bar' basis vectors on the 'plain'
#     #      basis vectors
#     #
#     # Calculate the rotation of the 'plain' stress vector to the
#     # 'bar' coordinate system given by the columns of the rotation matrix Rm.
#     #
#     # Example:
#     # The stress vector "stress" is given in the material coordinate
#     # system defined by the orientation matrix Rm. The following two
#     # transformations are equivalent:
#     #
#     #         t = stress_6v_to_3x3t (mat,stress);
#     #         t = (Rm*t*Rm');# in global coordinate system
#     #         t = (outputRm'*t*outputRm);# in output coordinate system
#     #         stress =stress_3x3t_to_6v (mat,t);# in output coordinate system
#     #
#     #        stress =mat.stress_vector_rotation(outputRm)...
#     #                  *mat.stress_vector_rotation(Rm')...
#     #                      *stress;# in output coordinate system
#
#     a11=Rm[1,1]; a12=Rm[1,2];
#     a21=Rm[2,1]; a22=Rm[2,2];
#     # switch self.reduction
#     #     case {'axisymm','strain'}
#     #         T =[[   a11^2,   a21^2, 0,         2*a11*a21]
#     #             [   a12^2,   a22^2, 0,         2*a12*a22]
#     #             [       0,       0, 1,                 0]
#     #             [ a11*a12, a21*a22, 0, a11*a22 + a12*a21]];
#     #             case 'stress'
#     T =[[   a11^2   a21^2       2*a11*a21]
#         [   a12^2   a22^2       2*a12*a22]
#         [ a11*a12 a21*a22 a11*a22+a12*a21]];
#     return T
# end
#
# function strainvectorrotation{MR<:DeforModelRed2DStress}(::Type{MR},
#   Rm::FMat{T})
#     # Calculate the rotation matrix for a strain vector.
#     #
#     #   Tbar = strain_vector_rotation(self,Rm)
#     #
#     # Rm = columns are components of 'bar' basis vectors on the 'plain'
#     #      basis vectors
#     #
#     # Calculate the rotation of the 'plain' strain vector to the
#     # 'bar' coordinate system given by the columns of the rotation matrix Rm.
#     #
#
#     a11=Rm[1,1]; a12=Rm[1,2];
#     a21=Rm[2,1]; a22=Rm[2,2];
#     # switch self.reduction
#     #     case {'axisymm','strain'}
#     #         Tbar =[ [     a11^2,     a21^2, 0,           a11*a21]
#     #                 [     a12^2,     a22^2, 0,           a12*a22]
#     #                 [         0,         0, 1,                 0]
#     #                 [ 2*a11*a12, 2*a21*a22, 0, a11*a22 + a12*a21]];
#     #   case 'stress'
#     Tbar =[ [     a11^2,     a21^2,           a11*a21]
#            [     a12^2,     a22^2,           a12*a22]
#            [ 2*a11*a12, 2*a21*a22, a11*a22 + a12*a21]];
#     return Tbar
# end
#
#
# ################################################################################
# # 2-D plane axially symmetric model
#
# function stressvectorrotation(::Type{MR},
#   Rm::FMat{T}) where {MR<:DeforModelRed2DAxisymm}
#
#     # Rm = columns are components of 'bar' basis vectors on the 'plain'
#     #      basis vectors
#     #
#     # Calculate the rotation of the 'plain' stress vector to the
#     # 'bar' coordinate system given by the columns of the rotation matrix Rm.
#     #
#     # Example:
#     # The stress vector "stress" is given in the material coordinate
#     # system defined by the orientation matrix Rm. The following two
#     # transformations are equivalent:
#     #
#     #         t = stress_6v_to_3x3t (mat,stress);
#     #         t = (Rm*t*Rm');# in global coordinate system
#     #         t = (outputRm'*t*outputRm);# in output coordinate system
#     #         stress =stress_3x3t_to_6v (mat,t);# in output coordinate system
#     #
#     #        stress =mat.stress_vector_rotation(outputRm)...
#     #                  *mat.stress_vector_rotation(Rm')...
#     #                      *stress;# in output coordinate system
#
#     a11=Rm[1,1]; a12=Rm[1,2];
#     a21=Rm[2,1]; a22=Rm[2,2];
#     T =[[   a11^2    a21^2  0          2*a11*a21]
#         [   a12^2    a22^2  0          2*a12*a22]
#         [       0        0  1                  0]
#         [ a11*a12  a21*a22  0  a11*a22 + a12*a21]];
#     return T
# end
#
# function strainvectorrotation(::Type{MR},
#   Rm::FMat{T}) where {MR<:DeforModelRed2DAxisymm}
#     # Calculate the rotation matrix for a strain vector.
#     #
#     #   Tbar = strain_vector_rotation(self,Rm)
#     #
#     # Rm = columns are components of 'bar' basis vectors on the 'plain'
#     #      basis vectors
#     #
#     # Calculate the rotation of the 'plain' strain vector to the
#     # 'bar' coordinate system given by the columns of the rotation matrix Rm.
#     #
#
#     a11=Rm[1,1]; a12=Rm[1,2];
#     a21=Rm[2,1]; a22=Rm[2,2];
#     Tbar =[ [     a11^2      a21^2  0            a11*a21]
#            [     a12^2      a22^2  0            a12*a22]
#            [         0          0  1                  0]
#            [ 2*a11*a12  2*a21*a22  0  a11*a22 + a12*a21]];
#     return Tbar
# end
#
# ################################################################################
# # Generic versions of rotations of stiffness and compliance matrices
#
# function rotatestiffness!(::Type{MR}, D::FMat{T},
#                                                            Rm::FMat{T}) where {MR<:DeforModelRed}
#     # Rotate constitutive stiffness matrix of the material.
#     #
#     #         function D=transform_stiffness(self,D,Rm)
#     #
#     # Rotate constitutive stiffness matrix of the material to the
#     # coordinate system given by the columns of the rotation matrix Rm.
#     T =stressvectorrotation(MR,Rm);
#     D = T*D*T';
#     return D
# end

#
# function rotatecompliance!(::Type{MR}, C::FMat{T},
#   Rm::FMat{T}) where {MR<:DeforModelRed}
#     # Rotate constitutive compliance matrix of the material.
#     #
#     #   C = rotate_compliance(self,C,Rm)
#     #
#     # Rotate constitutive compliance matrix of the material to the
#     # coordinate system given by the columns of the rotation matrix Rm.
#     Tbar =strainvectorrotation(MR,Rm);
#     C = Tbar*C*Tbar';
#     return C
# end
# function stressvectorrotation(::Type{DeforModelRed2DStrain}, Rm::FMat{T})
#
#     # Rm = columns are components of 'bar' basis vectors on the 'plain'
#     #      basis vectors
#     #
#     # Calculate the rotation of the 'plain' stress vector to the
#     # 'bar' coordinate system given by the columns of the rotation matrix Rm.
#     #
#     # Example:
#     # The stress vector "stress" is given in the material coordinate
#     # system defined by the orientation matrix Rm. The following two
#     # transformations are equivalent:
#     #
#     #         t = stress_6v_to_3x3t (mat,stress);
#     #         t = (Rm*t*Rm');# in global coordinate system
#     #         t = (outputRm'*t*outputRm);# in output coordinate system
#     #         stress =stress_3x3t_to_6v (mat,t);# in output coordinate system
#     #
#     #        stress =mat.stress_vector_rotation(outputRm)...
#     #                  *mat.stress_vector_rotation(Rm')...
#     #                      *stress;# in output coordinate system
#
#     a11=Rm[1,1]; a12=Rm[1,2];
#     a21=Rm[2,1]; a22=Rm[2,2];
#     # sigmax, sigmay, sigmaz, tauxy
#     T =[[   a11^2    a21^2  0          2*a11*a21]
#         [   a12^2    a22^2  0          2*a12*a22]
#         [       0        0  1                  0]
#         [ a11*a12  a21*a22  0  a11*a22 + a12*a21]];
#     return T
# end
#
# function strainvectorrotation(::Type{DeforModelRed2DStrain},  Rm::FMat{T})
#     # Calculate the rotation matrix for a strain vector.
#     #
#     #   Tbar = strain_vector_rotation(self,Rm)
#     #
#     # Rm = columns are components of 'bar' basis vectors on the 'plain'
#     #      basis vectors
#     #
#     # Calculate the rotation of the 'plain' strain vector to the
#     # 'bar' coordinate system given by the columns of the rotation matrix Rm.
#     #
#
#     a11=Rm[1,1]; a12=Rm[1,2];
#     a21=Rm[2,1]; a22=Rm[2,2];
#     Tbar =[ [     a11^2      a21^2  0            a11*a21]
#            [     a12^2      a22^2  0            a12*a22]
#            [         0          0  1                  0]
#            [ 2*a11*a12  2*a21*a22  0  a11*a22 + a12*a21]];
#     return Tbar
# end
# function strainvectorrotation(::Type{DeforModelRed3D}, Rm::FMat{T})
#     # Calculate the rotation matrix for a strain vector.
#     #
#     #   Tbar = strain_vector_rotation(self,Rm)
#     #
#     # Rm = columns are components of 'bar' basis vectors on the 'plain'
#     #      basis vectors
#     #
#     # Calculate the rotation of the 'plain' strain vector to the
#     # 'bar' coordinate system given by the columns of the rotation matrix Rm.
#     #
#
#     a11=Rm[1,1]; a12=Rm[1,2]; a13=Rm[1,3];
#     a21=Rm[2,1]; a22=Rm[2,2]; a23=Rm[2,3];
#     a31=Rm[3,1]; a32=Rm[3,2]; a33=Rm[3,3];
#     Tbar =[
#            [     a11^2     a21^2     a31^2           a11*a21           a11*a31           a21*a31]
#            [     a12^2     a22^2     a32^2           a12*a22           a12*a32           a22*a32]
#            [     a13^2     a23^2     a33^2           a13*a23           a13*a33           a23*a33]
#            [ 2*a11*a12 2*a21*a22 2*a31*a32 a11*a22 + a12*a21 a11*a32 + a12*a31 a21*a32 + a22*a31]
#            [ 2*a11*a13 2*a21*a23 2*a31*a33 a11*a23 + a13*a21 a11*a33 + a13*a31 a21*a33 + a23*a31]
#            [ 2*a12*a13 2*a22*a23 2*a32*a33 a12*a23 + a13*a22 a12*a33 + a13*a32 a22*a33 + a23*a32]];
#     return Tbar
# end


using Test, LinearAlgebra

@testset "/src/Degeneracy.jl" begin

using BellScenario: Degeneracy

@testset "symmetry_groups()" begin
    @testset "2x2 permutation group" begin
        transforms = [[1 0;0 1],[0 1;1 0]]
        elements = [[0 0],[1 0],[0 1],[1 1]]

        sym_groups = Degeneracy.symmetry_groups(transforms,elements)

        @test issetequal(elements, collect(Iterators.flatten(sym_groups)))
        @test length(sym_groups) == 3

        @test sym_groups[1] == [[0 0]]
        @test sym_groups[2] == [[1 0],[0 1]]
        @test sym_groups[3] == [[1 1]]
    end

    @testset "3x3 permutation group" begin
        transforms = Math.permutation_maps(3)
        elements = [
            [0 0 0],[0 0 1],[0 1 0],[0 1 1],
            [1 0 0],[1 0 1],[1 1 0],[1 1 1]
        ]

        sym_groups = Degeneracy.symmetry_groups(transforms,elements)

        @test issetequal(elements, collect(Iterators.flatten(sym_groups)))
        @test length(sym_groups) == 4

        @test sym_groups[1] == [[0 0 0]]
        @test sym_groups[2] == [[0 0 1],[0 1 0],[1 0 0]]
        @test sym_groups[3] == [[0 1 1],[1 0 1],[1 1 0]]
        @test sym_groups[4] == [[1 1 1]]
    end

    @testset "4x4 permutation group" begin
        transforms = Math.permutation_maps(4)
        elements = [
            [0 0 0 0],[0 0 0 1],[0 0 1 0],[0 0 1 1],
            [0 1 0 0],[0 1 0 1],[0 1 1 0],[0 1 1 1],
            [1 0 0 0],[1 0 0 1],[1 0 1 0],[1 0 1 1],
            [1 1 0 0],[1 1 0 1],[1 1 1 0],[1 1 1 1]
        ]

        sym_groups = Degeneracy.symmetry_groups(transforms,elements)

        @test issetequal(elements, collect(Iterators.flatten(sym_groups)))
        @test length(sym_groups) == 5

        @test sym_groups[1] == [[0 0 0 0]]
        @test sym_groups[2] == [[0 0 0 1],[0 0 1 0],[0 1 0 0],[1 0 0 0]]
        @test sym_groups[3] == [[0 0 1 1],[0 1 0 1],[0 1 1 0],[1 0 0 1],[1 0 1 0],[1 1 0 0]]
        @test sym_groups[4] == [[0 1 1 1],[1 0 1 1],[1 1 0 1],[1 1 1 0]]
        @test sym_groups[5] == [[1 1 1 1]]
    end
end

@testset "merge_symmetry_groups()" begin
    group1_sets = [[[1,1],[2,2]],[[3,3],[4,4],[5,5]],[[6,6]],[[7,7]]]
    group2_sets = [[[1,1],[3,3]],[[6,6],[7,7]],[[2,2],[4,4],[5,5]]]

    merged_groups = Degeneracy.merge_symmetry_groups(group1_sets,group2_sets)
    @test length(merged_groups) == 2
    @test all( # returned group sets do not overlap
        (x) -> x === nothing,
        indexin(merged_groups[1],merged_groups[2])
    )
    @test issetequal(merged_groups[1], [[1,1],[2,2],[3,3],[4,4],[5,5]])
    @test issetequal(merged_groups[2], [[6,6],[7,7]])
end

@testset "canonical_facets()" begin
    α_expt = (2,2)
    β_expt = (2,2)

    facets = [
        [0 0 0 0 0 -1 0 0 0],[0 0 0 0 0 0 -1 0 0],[0 0 0 0 0 0 0 -1 0],
        [0 0 0 0 0 0 0 0 -1],[0 -1 0 0 0 0 1 0 0],[0 -1 0 0 0 1 0 0 0],
        [0 0 -1 0 0 0 0 0 1],[0 0 -1 0 0 0 0 1 0],[0 0 0 -1 0 0 0 1 0],
        [0 0 0 -1 0 1 0 0 0],[0 0 0 0 -1 0 0 0 1],[0 0 0 0 -1 0 1 0 0],
        [0 -1 0 -1 0 1 1 1 -1],[0 -1 0 0 -1 1 1 -1 1],[0 0 -1 -1 0 1 -1 1 1],
        [0 0 -1 0 -1 -1 1 1 1],[-1 0 1 0 1 0 0 0 -1],[-1 0 1 1 0 0 0 -1 0],
        [-1 1 0 0 1 0 -1 0 0],[-1 1 0 1 0 -1 0 0 0],[-1 0 1 0 1 1 -1 -1 -1],
        [-1 0 1 1 0 -1 1 -1 -1],[-1 1 0 0 1 -1 -1 1 -1],[-1 1 0 1 0 -1 -1 -1 1]
    ]

    canonical_groups = Degeneracy.canonical_facets(α_expt, β_expt, facets, "no-signaling")

    @test issetequal(facets, collect(Iterators.flatten(canonical_groups)))
    @test length(canonical_groups) == 2
    @test issetequal(canonical_groups[1], [
        [0 0 0 0 0 -1 0 0 0],[0 0 0 0 0 0 -1 0 0],[0 0 0 0 0 0 0 -1 0],[0 0 0 0 0 0 0 0 -1],
        [0 -1 0 0 0 0 1 0 0],[0 -1 0 0 0 1 0 0 0],[0 0 -1 0 0 0 0 0 1],[0 0 -1 0 0 0 0 1 0],
        [0 0 0 -1 0 0 0 1 0],[0 0 0 -1 0 1 0 0 0],[0 0 0 0 -1 0 0 0 1],[0 0 0 0 -1 0 1 0 0],
        [-1 0 1 0 1 0 0 0 -1],[-1 0 1 1 0 0 0 -1 0],[-1 1 0 0 1 0 -1 0 0],[-1 1 0 1 0 -1 0 0 0]
    ])
    @test issetequal(canonical_groups[2], [
        [0 -1 0 -1 0 1 1 1 -1],[0 -1 0 0 -1 1 1 -1 1],[0 0 -1 -1 0 1 -1 1 1],[0 0 -1 0 -1 -1 1 1 1],
        [-1 0 1 0 1 1 -1 -1 -1],[-1 0 1 1 0 -1 1 -1 -1],[-1 1 0 0 1 -1 -1 1 -1],[-1 1 0 1 0 -1 -1 -1 1]
    ])
end

@testset "canonical_vertices()" begin
    @testset "chsh vertices" begin
        α1 = (2,2)
        β1 = (2,2)

        vertices = [
            [1;0;0;0;0;0;0;0;0],[1;0;0;0;1;0;0;0;0],[1;0;0;1;0;0;0;0;0],[1;0;0;1;1;0;0;0;0],
            [1;0;1;0;0;0;0;0;0],[1;0;1;0;1;0;0;0;1],[1;0;1;1;0;0;0;1;0],[1;0;1;1;1;0;0;1;1],
            [1;1;0;0;0;0;0;0;0],[1;1;0;0;1;0;1;0;0],[1;1;0;1;0;1;0;0;0],[1;1;0;1;1;1;1;0;0],
            [1;1;1;0;0;0;0;0;0],[1;1;1;0;1;0;1;0;1],[1;1;1;1;0;1;0;1;0],[1;1;1;1;1;1;1;1;1]
        ]

        canonical_vertices = Degeneracy.canonical_vertices(α1, β1, vertices, "no-signaling")

        @test length(canonical_vertices) == 1
        @test length(canonical_vertices[1]) == 16
        @test issetequal(canonical_vertices[1], vertices)
    end

    @testset "prepare and measure vertices" begin
        α = (3,1)
        β = (1,3)

        vertices = [
            [1 1 0 0 0 0 0]',[1 1 1 0 0 0 0]',[1 1 0 1 0 0 0]',[1 1 1 1 0 0 0]',[1 1 1 0 0 0 1]',[1 1 0 1 0 1 0]',
            [1 1 0 0 0 1 1]',[1 0 0 0 1 0 0]',[1 0 0 0 1 1 0]',[1 0 0 0 1 0 1]',[1 0 1 1 1 0 0]',[1 0 1 0 1 0 1]',
            [1 0 0 1 1 1 0]',[1 0 0 0 1 1 1]',[1 0 0 0 0 0 0]',[1 0 1 0 0 0 0]',[1 0 0 1 0 0 0]',[1 0 1 1 0 0 0]',
            [1 0 0 0 0 1 1]',[1 0 0 0 0 1 0]',[1 0 0 0 0 0 1]'
        ]

        canonical_vertices = Degeneracy.canonical_vertices(α, β, vertices, "normalized")

        @test length(canonical_vertices) == 2
        @test length(canonical_vertices[1]) == 3
        @test length(canonical_vertices[2]) == 18
        @test issetequal(canonical_vertices[1], [[1 1 1 1 0 0 0]',[1 0 0 0 0 0 0]',[1 0 0 0 1 1 1]'])
        @test issetequal( canonical_vertices[2],[
            [1 1 0 0 0 0 0]',[1 1 1 0 0 0 0]',[1 1 0 1 0 0 0]',[1 1 1 0 0 0 1]',[1 1 0 1 0 1 0]',
            [1 1 0 0 0 1 1]',[1 0 0 0 1 0 0]',[1 0 0 0 1 1 0]',[1 0 0 0 1 0 1]',[1 0 1 1 1 0 0]',
            [1 0 0 1 1 1 0]',[1 0 1 0 0 0 0]',[1 0 0 1 0 0 0]',[1 0 1 1 0 0 0]',[1 0 1 0 1 0 1]',
            [1 0 0 0 0 1 1]',[1 0 0 0 0 1 0]',[1 0 0 0 0 0 1]'
        ])
    end
end

@testset "bipartite_input_relabels()" begin

    @testset "no-signaling subspace" begin
        id = diagm(0 => fill(1, 9))
        flip_alice_input = [
            1 0 0 0 0 0 0 0 0;
            0 0 1 0 0 0 0 0 0;
            0 1 0 0 0 0 0 0 0;
            0 0 0 1 0 0 0 0 0;
            0 0 0 0 1 0 0 0 0;
            0 0 0 0 0 0 0 1 0;
            0 0 0 0 0 0 0 0 1;
            0 0 0 0 0 1 0 0 0;
            0 0 0 0 0 0 1 0 0;
        ];
        flip_bob_input = [
            1 0 0 0 0 0 0 0 0;
            0 1 0 0 0 0 0 0 0;
            0 0 1 0 0 0 0 0 0;
            0 0 0 0 1 0 0 0 0;
            0 0 0 1 0 0 0 0 0;
            0 0 0 0 0 0 1 0 0;
            0 0 0 0 0 1 0 0 0;
            0 0 0 0 0 0 0 0 1;
            0 0 0 0 0 0 0 1 0;
        ];
        flip_ab_input = [
            1 0 0 0 0 0 0 0 0;
            0 0 1 0 0 0 0 0 0;
            0 1 0 0 0 0 0 0 0;
            0 0 0 0 1 0 0 0 0;
            0 0 0 1 0 0 0 0 0;
            0 0 0 0 0 0 0 0 1;
            0 0 0 0 0 0 0 1 0;
            0 0 0 0 0 0 1 0 0;
            0 0 0 0 0 1 0 0 0;
        ]

        @test issetequal(
            Degeneracy.bipartite_input_relabels((2,2), (2,2), "no-signaling"),
            [id, flip_alice_input, flip_bob_input, flip_ab_input]
        )
    end

    test_cases = [
        ("no-signaling", 8),
        ("fixed-direction", 10),
        ("normalized", 12),
        ("generalized", 16)
    ]
    @testset "basic tests for subspace: $(case[1])" for case in test_cases
        maps = Degeneracy.bipartite_input_relabels((2,2),(2,2),case[1])

        dim = case[2]

        @test size(maps) == size(unique(maps))

        for i in 1:size(maps)[1]
            map = maps[i]

            @test size(map) == (dim+1,dim+1)
            @test map[1,1] == 1
            @test map[2:end,1] == zeros(Int64, dim)
            @test map[1,2:end] == zeros(Int64, dim)
        end
    end
end

@testset "bipartite_output_relabels()" begin
    @testset "dichotomic case" begin
        @testset "no-signaling" begin
            (α_relabels, β_relabels) = Degeneracy.bipartite_output_relabels((2,2),(2,2),"no-signaling")

            @test size(α_relabels) == (4,)
            @test size(β_relabels) == (4,)

            @test α_relabels[1] == diagm(0=>fill(1, 9))
            @test β_relabels[1] == diagm(0=>fill(1, 9))

            @test α_relabels[2] == [
                1 0 0 0 0 0  0 0 0;
                1 -1 0 0 0 0 0 0 0;
                0 0 1 0 0 0 0 0 0;
                0 0 0 1 0 0 0 0 0;
                0 0 0 0 1 0 0 0 0;
                0 0 0 1 0 -1 0 0 0;
                0 0 0 0 1 0 -1 0 0;
                0 0 0 0 0 0 0 1 0;
                0 0 0 0 0 0 0 0 1;
            ]
            @test α_relabels[3] == [
                 1 0 0 0 0 0 0 0 0;
                 0 1 0 0 0 0 0 0 0;
                 1 0 -1 0 0 0 0 0 0;
                 0 0 0 1 0 0 0 0 0;
                 0 0 0 0 1 0 0 0 0;
                 0 0 0 0 0 1 0 0 0;
                 0 0 0 0 0 0 1 0 0;
                 0 0 0 1 0 0 0 -1 0;
                 0 0 0 0 1 0 0 0 -1;
             ]
             @test β_relabels[2] == [
                1 0 0 0 0 0 0 0 0;
                0 1 0 0 0 0 0 0 0;
                0 0 1 0 0 0 0 0 0;
                1 0 0 -1 0 0 0 0 0;
                0 0 0 0 1 0 0 0 0;
                0 1 0 0 0 -1 0 0 0;
                0 0 0 0 0 0 1 0 0;
                0 0 1 0 0 0 0 -1 0;
                0 0 0 0 0 0 0 0 1;
            ]
            @test β_relabels[3] == [
                1 0 0 0 0 0 0 0 0;
                0 1 0 0 0 0 0 0 0;
                0 0 1 0 0 0 0 0 0;
                0 0 0 1 0 0 0 0 0;
                1 0 0 0 -1 0 0 0 0;
                0 0 0 0 0 1 0 0 0;
                0 1 0 0 0 0 -1 0 0;
                0 0 0 0 0 0 0 1 0;
                0 0 1 0 0 0 0 0 -1;
            ]
            @test β_relabels[4] == β_relabels[2]*β_relabels[3]
        end
    end

    @testset "relabel set contains the inverse of each element" begin
        test_cases = [
            ((2,2),(2,2)),
            ((3,2),(2,2)),
            ((2,3),(2,3)),
            ((3,1),(1,3))
        ]

        @testset "generalized $(case[1]), $(case[2])" for case in test_cases
            (α_relabels, β_relabels) = Degeneracy.bipartite_output_relabels(case[1],case[2],"generalized")

            dim = Behavior.dimension(case[1],case[2],"generalized")

            # in generalized coordinates, output relabelings inverted through
            # their own transpose.
            for relabel in α_relabels
                parity = det(relabel)
                @test (parity == 1) | (parity == -1)
                @test relabel*relabel' == diagm( 0 => fill(1,dim))
            end

            for relabel in β_relabels
                parity = det(relabel)
                @test (parity == 1) | (parity == -1)
                @test relabel*relabel' == diagm( 0 => fill(1,dim))
            end
        end

        @testset "no-signaling $(case[1]), $(case[2])" for case in test_cases
            (α_relabels, β_relabels) = Degeneracy.bipartite_output_relabels(case[1],case[2],"no-signaling")

            dim = Behavior.dimension(case[1],case[2],"no-signaling")
            Id = diagm(0 => fill(1,dim))

            foreach( relabel_test->begin
                matches = filter(relabel_inv->relabel_test*relabel_inv == Id, α_relabels)
                @test (1,) == size(matches)
            end, α_relabels)

            foreach( relabel_test->begin
                matches = filter(relabel_inv->relabel_test*relabel_inv == Id, β_relabels)
                @test (1,) == size(matches)
            end, β_relabels)
        end

        @testset "fixed-direction $(case[1]), $(case[2])" for case in test_cases
            (α_relabels, β_relabels) = Degeneracy.bipartite_output_relabels(case[1],case[2],"fixed-direction")

            dim = Behavior.dimension(case[1],case[2],"fixed-direction")
            Id = diagm(0 => fill(1,dim))

            # check that there is exactly 1 inverting map for every map in the
            # relabeling set
            foreach( relabel_test->begin
                matches = filter(relabel_inv->relabel_test*relabel_inv == Id, α_relabels)
                @test (1,) == size(matches)
            end, α_relabels)

            foreach( relabel_test->begin
                matches = filter(relabel_inv->relabel_test*relabel_inv == Id, β_relabels)
                @test (1,) == size(matches)
            end, β_relabels)
        end
    end
end

end


# Defines the Entry type and associated methods. Entry Represents
# a single entry in Pascal's triangle - a location (n,k) and and
# associated value = binomial(n,k).

"""
    Entry(n, k, val)
    Entry(n, k)
    Entry(n => k)
    Entry{V}(n, k)
    Entry(e::Entry)

Represents a single entry in Pascal's triangle. The supplied `val` is not checked for
correctness but arguments that only specify `n` and `k` must satisfy 0 ≤ k ≤ n and must
both be integer types. The argument `V` specifies a type for the value field.

# Examples
```jldoctest
julia> Entry(6,2)
Entry{Int64,Int64}(6, 2, 15)

julia> Entry{Float64}(6, 2)
Entry{Int64,Float64}(6, 2, 15.0)

julia> Entry(6,2,10) # incorrect value not checked
Entry{Int64,Int64}(6, 2, 10)
```
"""
mutable struct Entry{V <: Real}
    n::Int
    k::Int
    val::V
    function Entry{V}(n,k,val) where {V <: Real}
        0 ≤ k ≤ n || throw(ArgumentError("Entry requires 0 ≤ k ≤ n but n:$n, k:$k"))
        return new{V}(n,k,val)
    end
end
Entry(n,k,val) = Entry{typeof(val)}(n,k,val)
Entry(e::Entry) = Entry(e.n, e.k, e.val)
Entry{V}(n,k) where {V <: Real} = Entry{V}(n, k, binomial(promote(n, k)...))
Entry(n,k) = Entry{Int}(n,k)
Entry(p::Pair) = Entry(first(p), last(p))

rownumber(e::Entry) = e.n
rowposition(e::Entry) = e.k
value(e::Entry) = e.val

Base.show(io::IO, ::MIME"text/plain", e::Entry) = print(io, "$(typeof(e).name)($(e.n), $(e.k), $(e.val))")
Base.show(io::IO, e::Entry) = print(io, "($(e.n), $(e.k), $(e.val))")

"""
    ==(a::Entry, b::Entry)

Two entries are equal if they represent the same entry in Pascal's triangle. That is,
they are from the same row, row position and have the same value, although the types
need not match.
"""
Base.:(==)(a::Entry, b::Entry) = a.n == b.n && a.k == b.k && a.val == b.val

"""
    isless(a::Entry, b::Entry)

Entries are ordered only by their value, ignoring their rows and row positions.
"""
Base.isless(a::Entry, b::Entry) = a.val < b.val

"""
    ≈(a::Entry, b::Entry)

Entries are approximately equal if they represent the same location within Pascal's
triangle and their values are approximately equal.
"""
Base.isapprox(a::Entry, b::Entry) = a.n == b.n && a.k == b.k && isapprox(a.val, b.val)

"""
    +(a::Entry, b::Entry)

Adding adjacent entries gives the entry directly beneath the two.
"""
function Base.:+(a::Entry, b::Entry)
    isadjacent(a,b) || throw(NonAdjacentError())
    return Entry(a.n+1, max(a.k, b.k), a.val + b.val)
end

"""
    -(a::Entry, b::Entry)

Subtracting two entries that are directly above each other gives the entry directly
adjacent to the upper one, utilising the well-known relation on Pascal's triangle.
"""
function Base.:-(a::Entry, b::Entry)
    issubtractable(a,b) || throw(NonAdjacentError())
    return Entry(b.n, 2a.k - b.k - 1, a.val - b.val)
end

isfirst(a::Entry) = a.n == 0 && a.k == 0
isatleft(a::Entry) = a.k ≤ 0
isatright(a::Entry) = a.k ≥ a.n
isadjacent(a::Entry, b::Entry) = a.n == b.n && abs(a.k - b.k) == 1
isinterior(a::Entry) = a.n ≥ 2 && a.k ≥ 1 && a.k < a.n
issubtractable(a::Entry, b::Entry) = a.n == b.n + 1 && isinterior(a) && 0 ≤ a.k - b.k ≤ 1
isvalid(a::Entry) = 0 ≤ a.n && 0 ≤ a.k ≤ a.n && binomial(a.n, a.k) == a.val


using RCall, DataFrames

"""Load characters data"""
function initialize(label_a::Int, label_b::Int)::Tuple{DataFrames.DataFrame,DataFrames.DataFrame}
    @rput label_a
    @rput label_b
    R"""
    load('./mnist_train.rda')
    mnist = subset(mnist_train, mnist_train[,1] %in% c(label_a, label_b))
    mnist_means = aggregate(mnist[,-1], by = list(mnist[,1]), FUN = mean)
    """
    @rget mnist
    @rget mnist_means
    return mnist, mnist_means
end

"""Single pixel classifier based on mean pixel value"""
function baseClassifier_j(mnist_means::DataFrames.DataFrame, pixel::Int, j::Int, correct::Int)
    if abs(pixel - mnist_means[1,j]) <= abs(pixel - mnist_means[2,j])
        label = mnist_means[1,1]
    else
        label = mnist_means[2,1]
    end
    return label, correct == label
end

"""Return which label is the most common in a vector"""
function findMostCommon(vector::Vector{Int}, a::Int, b::Int)
    count_a = 0; count_b = 0
    for el in vector
        if el == a
            count_a += 1
        else
            count_b += 1
        end
    end
    if count_a > count_b
        return a
    else
        return b
    end
end

function main(label_a::Int, label_b::Int, ϵ::Float64, γ::Float64)
    mnist, mnist_means = initialize(label_a, label_b)

    d = size(mnist, 2) - 1
    n = size(mnist, 1)
    w = ones(n)
    η = γ
    T = convert(Int, ceil(2/γ^2 * log(1/ϵ)))
    info("T = ", T)
    corrects = falses(n)
    labels = Vector{Int}(n)
    labels_per_classifier = Matrix{Int}(n, T)
    correct_final_label = mnist[:X5]
    final_label = Vector{Int64}(n);
    hits = falses(n)
    percent_hits = zeros(T)
    count_base_classifiers = zeros(d-1)

    for t in 1:T
        j = 1
        S = 0.0
        while S < 0.5 + γ
            p = w / sum(w)
            S = 0.0
            for i in 1:n
                label, correct_label = baseClassifier_j(mnist_means, Int64(mnist[i, j + 1]), j + 1, Int64(mnist[i, 1]))
                labels[i] = label
                corrects[i] = correct_label
            end

            S = sum(p .* corrects)

            j += 1
            if j > d
                warn("No valid weak learner found for iteration ", t)
            end
        end
        count_base_classifiers[j] = count_base_classifiers[j] + 1
        labels_per_classifier[:, t] = labels
        # Found weak learner corresponding to j-th baseClassifier
        # Update weights
        for i in 1:n
            w[i] = w[i] * exp(η * !corrects[i]) # Increase weights for those that were wrongly classified
        end
        p = w / sum(w)
        for i in 1:n
            final_label[i] = findMostCommon(labels_per_classifier[i,1:t], label_a, label_b)
            hits[i] = final_label[i] == correct_final_label[i]
        end

        percent_hits[t] = sum(hits)/length(hits)

        if t % 10 == 0
            info("t = ", t, " of ", T)
            info("Correct ratio: ", percent_hits[t])
        end
        if percent_hits[t] > 1-ϵ
          return percent_hits, count_base_classifiers
        end
    end
    return percentual_hits, count_base_classifiers
end


module MapsMocks

using CellwiseValues

export MockMap
export TestMap

import CellwiseValues: evaluate!
import CellwiseValues: return_size

struct MockMap{P} <: Map{P,1,P,1}
  val::P
end

function evaluate!(
  this::MockMap{P}, points::AbstractVector{P}, v::AbstractVector{P}) where P
  for (i,qi) in enumerate(points)
    v[i] = qi+this.val
  end
end

return_size(::MockMap, psize::Tuple{Int}) = psize

struct TestMap{P} <: Map{P,1,P,2}
  val::P
  dim::Int
end

function evaluate!(
  this::TestMap{P}, points::AbstractVector{P}, v::AbstractMatrix{P}) where P
  for j in 1:this.dim
    for (i,qi) in enumerate(points)
      v[j,i] = j*qi+this.val
    end
  end
end

return_size(this::TestMap, psize::Tuple{Int}) = (this.dim, psize[1])

end # module


using Pumas, PumasTutorials, CSV, Plots

data = PumasTutorials.tutorial_data("data/nca","SD_crossover_IV_oral")
data = CSV.read(data,missingstring="NA")
first(data,10)

timeu = u"hr"
concu = u"mg/L"
amtu  = u"mg"

pop = read_nca(data, id=:ID, time=:time, conc=:DV, amt=:DOSE, ii=24timeu,
    route=:Formulation, occasion=:OCC,timeu=timeu, concu=concu, amtu=amtu)

pop = read_nca(data, id=:ID, time=:time, conc=:DV, amt=:DOSE, ii=24timeu,
    route=:Formulation, occasion=:OCC,timeu=timeu, concu=concu, amtu=amtu, lloq=0.4concu)

NCA.auc(pop,auctype=:last,method=:linear)

NCA.auc(pop,auctype=:inf,method=:linuplogdown)

NCA.auc(pop[1], interval=(0,12).*timeu)

NCA.auc(pop[1], interval=[(0,12).*timeu,(0,6).*timeu])

NCA.auc_extrap_percent(pop)

NCA.lambdaz(pop)

NCA.lambdaz(pop, threshold=3)

NCA.lambdaz(pop, idxs=[18,19,20])

NCA.lambdaz(pop, slopetimes=[18.5,19,19.5].*timeu)

cmax = NCA.cmax(pop[1])

pop[1]

normalizedose(cmax,pop[1])

AUClast = NCA.auc(pop[1],auctype=:last)
normalizedose(AUClast,pop[1])

NCA.lambdazr2(pop)
NCA.lambdazadjr2(pop)
NCA.lambdazintercept(pop)
NCA.lambdaztimefirst(pop)
NCA.lambdaznpoints(pop)

NCA.tmax(pop)
NCA.cmin(pop)
NCA.tmin(pop)

NCA.tlast(pop)
NCA.clast(pop)

NCA.aumc(pop)
NCA.aumclast(pop)

NCA.thalf(pop)

NCA.cl(pop,ithdose=1)

NCA.vss(pop)
NCA.vz(pop)

NCA.accumulationindex(pop)

pop[2]
NCA.cl(pop[2])

report = NCAReport(pop,ithdose=1)
report = NCA.to_dataframe(report)

report = NCAReport(pop,ithdose=1,pred=true)
report = NCA.to_dataframe(report)

CSV.write("./tutorials/nca/report_SD_crossover_IV_oral.csv", report)

using PuMaSTutorials
PuMaSTutorials.tutorial_footer(WEAVE_ARGS[:folder],WEAVE_ARGS[:file])


using StatsBase, Turing
using Gadfly

# include("ASCIIPlot.jl")
import Gadfly.ElementOrFunction

# First add a method to the basic Gadfly.plot function for QQPair types (generated by Distributions.qqbuild())
Gadfly.plot(qq::QQPair, elements::ElementOrFunction...) = Gadfly.plot(x=qq.qx, y=qq.qy, Geom.point, Theme(highlight_width=0px), elements...)

# Now some shorthand functions
qqplot(x, y, elements::ElementOrFunction...) = Gadfly.plot(qqbuild(x, y), elements...)
qqnorm(x, elements::ElementOrFunction...) = qqplot(Normal(), x, Guide.xlabel("Theoretical Normal quantiles"), Guide.ylabel("Observed quantiles"), elements...)

NSamples = 5000

@model gdemo_fw() = begin
  s ~ InverseGamma(2,3)
  m ~ Normal(0, sqrt(s))
  y ~ MvNormal([m; m; m], [sqrt(s) 0 0; 0 sqrt(s) 0; 0 0 sqrt(s)])
end

@model gdemo_bk(x) = begin
  # Backward Step 1: theta ~ theta | x
  s ~ InverseGamma(2,3)
  m ~ Normal(0,sqrt(s))
  x ~ MvNormal([m; m; m], [sqrt(s) 0 0; 0 sqrt(s) 0; 0 0 sqrt(s)])
  # Backward Step 2: x ~ x | theta
  y ~ MvNormal([m; m; m], [sqrt(s) 0 0; 0 sqrt(s) 0; 0 0 sqrt(s)])
end

fw = HMCDA(NSamples, 0.9, 0.1)
# bk = Gibbs(10, PG(10,10, :s, :y), HMC(1, 0.25, 5, :m));
bk = HMCDA(50, 0.9, 0.1);

s = sample(gdemo_fw(), fw);
# describe(s)

N = div(NSamples, 50)

x = [s[:y][1]...]
s_bk = Array{Turing.Chain}(undef, N)

simple_logger = Base.CoreLogging.SimpleLogger(stderr, Base.CoreLogging.Debug)
with_logger(simple_logger) do
  i = 1
  while i <= N
    try
      s_bk[i] = sample(gdemo_bk(x), bk);
      x = [s_bk[i][:y][end]...];
      i += 1
    catch
    end
  end
do

s2 = vcat(s_bk...);
# describe(s2)

qqplot(s[:m], s2[:m])
qqplot(s[:s], s2[:s])

using Test

qqm = qqbuild(s[:m], s2[:m])
X = qqm.qx
y = qqm.qy
slope = (1 / (transpose(X) * X)[1] * transpose(X) * y)[1]
@test slope ≈ 1.0 atol=0.1

# NOTE: test for s is not stable
#       probably due to the transformation
# qqs = qqbuild(s[:s], s2[:s])
# X = qqs.qx
# y = qqs.qy
# slope = (1 / (transpose(X) * X)[1] * transpose(X) * y)[1]
# @test slope ≈ 1.0 atol=0.1


# ---
# title: 867. Transpose Matrix
# id: problem867
# author: Qling
# date: 2021-02-25
# difficulty: Easy
# categories: Array
# link: <https://leetcode.com/problems/transpose-matrix/description/>
# hidden: true
# ---
#
# Given a matrix `A`, return the transpose of `A`.
#
# The transpose of a matrix is the matrix flipped over it's main diagonal,
# switching the row and column indices of the matrix.
#
#
# ![](https://assets.leetcode.com/uploads/2019/10/20/hint_transpose.png)
#
#
#
# **Example 1:**
#
#
#
#     Input: [[1,2,3],[4,5,6],[7,8,9]]
#     Output: [[1,4,7],[2,5,8],[3,6,9]]
#
#
# **Example 2:**
#
#
#
#     Input: [[1,2,3],[4,5,6]]
#     Output: [[1,4],[2,5],[3,6]]
#
#
#
#
# **Note:**
#
#   1. `1 <= A.length <= 1000`
#   2. `1 <= A[0].length <= 1000`
#
#
## @lc code=start
using LeetCode

function transpose_matrix(matrix::Vector{Vector{Int}})::Vector{Vector{Int}}
    m, n = length(matrix), length(matrix[1])

    transpose_mat = [[0 for _ = 1:m] for _ = 1:n]

    for i = 1:m, j = 1:n
        transpose_mat[j][i] = matrix[i][j]
    end

    return transpose_mat
end
## @lc code=end


#
# Do work here ...
#

# This is a 1x3 Matrix{Float64}
M = [1.1 2.3 4.5]

# This is a 3-element Vector{Float64}
v = [1.1; 2.3; 4.5]
println("Julia side: I will return a Vector{Float64}: ", v)

return v


module Windows
using ..Util
import SpecialFunctions: besseli
using LinearAlgebra: Diagonal, SymTridiagonal, eigen!, mul!, rmul!
using FFTW

export  rect,
        hanning,
        hamming,
        tukey,
        cosine,
        lanczos,
        triang,
        bartlett,
        gaussian,
        bartlett_hann,
        blackman,
        kaiser,
        dpss,
        dpsseig

# this gets shared between the docstrings of all the window functions because
# they share these args.
const zerophase_docs = """
If `zerophase` is `false` (the default) the window is centered around index
`(n+1)/2`, which is commonly used for FIR filter design. These are often used
for FIR filter design, and usually odd-length. Note that for even-length windows
this will cause the window center to be located between samples.

If `zerophase` is `true` the window is centered around index 1 (with the
negative half wrapped to the end of the vector). Additionally this creates a
"periodic" window, which means that if there is no padding then the left and
right endpoints of the window wrap around to the same sample, so the window
length is the same as an `n+1`-length non-`zerophase` window. Alternatively you
can think of the continuous `zerophase` window being of width `n` and the
non-`zerophase` window being of length `n-1`. `zerophase` windows are often used
in FFT processing, and are usually even-length.
"""

function twoD_docs(arg=nothing)
    argstr = if arg !== nothing
        "`$arg`, "
    else
        ""
    end
    """
    Providing a `dims` `Tuple` rather than a single `n` constructs a 2D window.
    $argstr`padding` and `zerophase` can then be given either as a single value
    for both horizontal and vertical or a 2-tuple to specify them separately.
    """
end

"""
    padplot(plotstr)

Takes a multiline string and pre-pads it so that it shows up as
preformatted code when included in a docstring.
"""
function padplot(plotstr)
    minpad = typemax(Int)
    lines = split(plotstr, "\n")
    # count the minimum number of spaces preceeding any line
    for line in lines
        line == "" && continue
        minpad = min(minpad, length(match(r"^ *", line).match))
    end

    pad = " " ^ (4-minpad)
    join((string(pad, line) for line in lines), "\n")
end

include("winplots.jl")

"""
    makewindow(winfunc::Function, n::Integer, padding::Integer, zerophase::Bool)

Generate a discrete window vector of the given length with `padding` zeros.
`winfunc` should be a function giving the window value in the range of
[-0.5, 0.5]. The function is assumed to be 0 outside of this range.

$zerophase_docs

Example:

```
function hanning(n::Integer; padding::Integer=0, zerophase::Bool=false)
    makewindow(n, padding, zerophase) do x
        0.5*(1+cos(2pi*x))
    end
end
```
"""
function makewindow(winfunc::Function, n::Integer, padding::Integer, zerophase::Bool)
    if n < 0
        throw(ArgumentError("`n` must be nonnegative"))
    end
    if padding < 0
        throw(ArgumentError("`padding` must be nonnegative"))
    end
    win = zeros(n+padding)
    if n == 1
        win[1] = winfunc(0.0)
    elseif zerophase
        # note that the endpoint of the window gets set in both lines. In the
        # unpadded case this will set the same index (which shouldn't make a
        # difference if the window is symmetric), but it's necessary for when
        # there's padding, which ends up in the center of the vector length
        # n÷2+1
        win[1:n÷2+1] .= winfunc.(range(0.0, stop=(n÷2)/n, length=n÷2+1))
        # length n÷2
        win[end-n÷2+1:end] .= winfunc.(range(-(n÷2)/n, stop=-1/n, length=n÷2))
    else
        win[1:n] .= winfunc.(range(-0.5, stop=0.5, length=n))
    end

    win
end

#
# Window functions
#

# equations rendered into Unicode with https://arthursonzogni.com/Diagon/

"""
$rect_winplot

    rect(n::Integer; padding::Integer=0, zerophase::Bool=false)
    rect(dims; padding=0, zerophase=false)

Rectangular window of length `n`, padded with `padding` zeros. This window is 1
within the window, and 0 outside of it.

$(twoD_docs())

$zerophase_docs
"""
function rect(n::Integer; padding::Integer=0, zerophase::Bool=false)
    makewindow(n, padding, zerophase) do x
        1.0
    end
end

"""
$hanning_winplot

    hanning(n::Integer; padding::Integer=0, zerophase::Bool=false)
    hanning(dims; padding=0, zerophase=false)

Hanning window of length `n` with `padding` zeros. The Hanning (or Hann) window
is a raised-cosine window that reaches zero at the endpoints.

The window is defined by sampling the continuous function:

           1 + cos(2πx)
    w(x) = ──────────── = cos²(πx)
                2

in the range `[-0.5, 0.5]`

The `hanning` window satisfies the Constant Overlap-Add (COLA) property with an
hop of 0.5, which means that adding together a sequence of delayed windows with
50% overlap will result in a constant signal. This is useful when synthesizing
a signal from a number of overlapping frames (each with a roughly rectangular
window), to eliminate windowing amplitude modulation.

Note that the `hanning` window is the `cosine` window squared.

$(twoD_docs())

$zerophase_docs
"""
function hanning(n::Integer; padding::Integer=0, zerophase::Bool=false)
    makewindow(n, padding, zerophase) do x
        0.5*(1+cos(2pi*x))
    end
end

"""
$hamming_winplot

    hamming(n::Integer; padding::Integer=0, zerophase::Bool=false)
    hamming(dims; padding=0, zerophase=false)

Hamming window of length `n` with `padding` zeros. The Hamming window does not
reach zero at the endpoints and so has a shallower frequency roll-off when
compared to the Hanning window, but is designed to cancel the first side-lobe.

The window is defined by sampling the continuous function:

    w(x) = 0.54 + 0.46*cos(2pi*x)

in the range `[-0.5, 0.5]`

$(twoD_docs())

$zerophase_docs
"""
function hamming(n::Integer; padding::Integer=0, zerophase::Bool=false)
    makewindow(n, padding, zerophase) do x
        0.54 + 0.46*cos(2pi*x)
    end
end

"""
$tukey_winplot

    tukey(n::Integer, α::Real; padding::Integer=0, zerophase::Bool=false)
    tukey(dims, α; padding=0, zerophase=false)

Tukey window of length `n` with `padding` zeros. The Tukey window has a flat top
and reaches zero at the endpoints, with a sinusoidal transition area
parameterized by `α`. For `α == 0`, the window is equivalent to a rectangular
window. For `α == 1`, the window is a Hann window.

The window is defined by sampling the continuous function:

           ⎛              ⎛    ⎛    1 - α⎞⎞
           ⎜      1 + cos ⎜2πα ⎜x + ─────⎟⎟             1 - α
           ⎜              ⎝    ⎝      2  ⎠⎠         x ≤ ─────
           ⎜      ─────────────────────────               2
           ⎜                  2
           ⎜
    w(x) = ⎜      1                                 -α/2 < x ≤ α/2
           ⎜
           ⎜              ⎛    ⎛    1 - α⎞⎞
           ⎜      1 + cos ⎜2πα ⎜x - ─────⎟⎟             1 - α
           ⎜              ⎝    ⎝      2  ⎠⎠         x > ─────
           ⎜      ─────────────────────────               2
           ⎝                  2

in the range `[-0.5, 0.5]`

$(twoD_docs("α"))

$zerophase_docs
"""
function tukey(n::Integer, α::Real; padding::Integer=0, zerophase::Bool=false)
    # check that α is reasonable
    !(0 <= α <= 1) && error("α must be in the range 0 <= α <= 1.")

    # if α is less than machine precision, call it zero and return the
    # rectangular window for this length.  if we don't short circuit this
    # here, it will blow up below.
    abs(α) <= eps() && return rect(n; padding=padding, zerophase=zerophase)

    makewindow(n, padding, zerophase) do x
        if x <= -(1-α)/2
            0.5*(1 + cos(2pi/α*(x+(1-α)/2)))
        elseif x <= (1-α)/2
            1.0
        else
            0.5*(1 + cos(2pi/α*(x-(1-α)/2)))
        end
    end
end

"""
$cosine_winplot

    cosine(n::Integer; padding::Integer=0, zerophase::Bool=false)
    cosine(dims; padding=0, zerophase=false)

Cosine window of length `n` with `padding` zeros. The cosine window is the first
lobe of a cosine function (with the zero crossings at +/- π as endpoints). Also
called the sine window.

The window is defined by sampling the continuous function:

    w(x) = cos(πx)

in the range `[-0.5, 0.5]`

Note that the cosine window is the square root of the `hanning` window, so it is
sometimes used when you are applying the window twice, such as the analysis and
synthesis steps of an STFT.

$(twoD_docs())

$zerophase_docs
"""
function cosine(n::Integer; padding::Integer=0, zerophase::Bool=false)
    makewindow(n, padding, zerophase) do x
        cos(pi*x)
    end
end

"""
$lanczos_winplot

    lanczos(n::Integer; padding::Integer=0, zerophase::Bool=false)
    lanczos(dims; padding=0, zerophase=false)

Lanczos window of length `n` with `padding` zeros. The Lanczos window is the
main lobe of a `sinc` function.

The window is defined by sampling the continuous function:

                      sin(2πx)
    w(x) = sinc(2x) = ────────
                         2πx

in the range `[-0.5, 0.5]`

$(twoD_docs())

$zerophase_docs
"""
function lanczos(n::Integer; padding::Integer=0, zerophase::Bool=false)
    makewindow(n, padding, zerophase) do x
        sinc(2x)
    end
end

"""
$triang_winplot

    triang(n::Integer; padding::Integer=0, zerophase::Bool=false)
    triang(dims; padding=0, zerophase=false)

Triangular window of length `n` with `padding` zeros. The Triangular window does
not reach zero at the endpoints. For odd `n` the `triang` window is the center
`n` points of an `n+2`-point [`bartlett`](@ref) window (i.e. the samples just
outside the window would be zero). For even `n` the window slope is the same as
the `n-1` window but delayed by a half sample so the zero points would be 1/2
sample past the ends of the window.

The window is defined by sampling the continuous function:

            ⎛    2(n-1)
            ⎜1 - ────── abs(x)     n is even
            ⎜       n
    w(x) =  ⎜
            ⎜    2(n-1)
            ⎜1 - ────── abs(x)     n is odd
            ⎝     n+1

in the range `[-0.5, 0.5]`.

$(twoD_docs())

$zerophase_docs

When `zerophase` is `true` substitute `n+1` for `n` in the above window
expressions.
"""
function triang(n::Integer; padding::Integer=0, zerophase::Bool=false)
    # for the purpose of calculating the slope of the window, consider `n` to be
    # 1 larger to compensate for the fact that `zerophase` gives a periodic
    # window
    m = zerophase ? n+1 : n
    scale = iseven(m) ? 2(m-1)/m : 2(m-1)/(m+1)
    makewindow(n, padding, zerophase) do x
        1 - scale*abs(x)
    end
end

"""
$bartlett_winplot

    bartlett(n::Integer; padding::Integer=0, zerophase::Bool=false)
    bartlett(dims; padding=0, zerophase=false)

Bartlett window of length `n`. The Bartlett window is a triangular window that
reaches 0 at the endpoints. This is equivalent to convolving two rectangular
windows of length `(n-1)/2` and adding the zero endpoints. See `triang` for a
window that does not reach zero at the endpoints.

The window is defined by sampling the continuous function:

    1 - abs(2x)

in the range `[-0.5, 0.5]`

$(twoD_docs())

$zerophase_docs
"""
function bartlett(n::Integer; padding::Integer=0, zerophase::Bool=false)
    makewindow(n, padding, zerophase) do x
        1 - abs(2x)
    end
end

"""
$gaussian_winplot

    gaussian(n::Integer, σ::Real; padding::Integer=0, zerophase::Bool=false)
    gaussian(dims, σ; padding=0, zerophase=false)

Gives an n-sample gaussian window defined by sampling the function:

            ⎛        2⎞
            ⎜-1   ⎛x⎞ ⎟
            ⎜── ⋅ ⎜─⎟ ⎟
            ⎝ 2   ⎝σ⎠ ⎠
    w(x) = e

in the range `[-0.5,0.5]`. This means that for `σ=0.5` the endpoints of the
window will correspond to 1 standard deviation away from the center.

$(twoD_docs("σ"))

$zerophase_docs
"""
function gaussian(n::Integer, σ::Real; padding::Integer=0, zerophase::Bool=false)
    σ > 0.0 || error("σ must be positive")
    makewindow(n, padding, zerophase) do x
        exp(-0.5*(x/σ)^2)
    end
end

"""
$bartlett_hann_winplot

    bartlett_hann(n::Integer; padding::Integer=0, zerophase::Bool=false)
    bartlett_hann(dims; padding=0, zerophase=false)

Bartlett-Hann window of length `n` with `padding` zeros. The Bartlett-Hann
window is a weighted sum of the Bartlett and Hann windows.

The window is defined by sampling the continuous function:

    w(x) = 0.62 - 0.48*abs(x) + 0.38*cos(2π*x)

in the range `[-0.5, 0.5]`

$(twoD_docs())

$zerophase_docs
"""
function bartlett_hann(n::Integer; padding::Integer=0, zerophase::Bool=false)
    a0, a1, a2 = 0.62, 0.48, 0.38
    makewindow(n, padding, zerophase) do x
        a0 - a1*abs(x) + a2*cos(2pi*x)
    end
end

"""
$blackman_winplot

    blackman(n::Integer; padding::Integer=0, zerophase::Bool=false)
    blackman(dims; padding=0, zerophase=false)

Approximates the "Exact" Blackman window. This is the generalized Blackman
window with α = 0.16.

The window is defined by sampling the continuous function:

    w(x) = 0.42 + 0.5*cos(2π*x) + 0.08*cos(4π*x)

in the range `[-0.5, 0.5]`

$(twoD_docs())

$zerophase_docs
"""
function blackman(n::Integer; padding::Integer=0, zerophase::Bool=false)
    a0, a1, a2 = 0.42, 0.5, 0.08
    makewindow(n, padding, zerophase) do x
        a0 + a1*cos(2pi*x) + a2*cos(4pi*x)
    end
end

"""
$kaiser_winplot

    kaiser(n::Integer, α::Real; padding::Integer=0, zerophase::Bool=false)
    kaiser(dims, α; padding=0, zerophase=false)

Kaiser window of length `n` parameterized by `α`. The Kaiser window approximates
the DPSS window (given by `dpss`), using a simplified definition relying on a
Bessel function. Larger values for `α` give a wider main lobe but have lower
sidelobes. Typically `α` is set around 3.

The window is defined by sampling the continuous function:

             ⎛  ⎛   _________⎞⎞
             ⎜  ⎜  ╱        2⎟⎟
    w(x) = I₀⎝πα⎝╲╱ 1 - (2x) ⎠⎠
           ────────────────────
                   I₀(πα)

in the range `[-0.5, 0.5]`

Where I₀(⋅) is the zeroth-order modified Bessel function of the first kind.

$(twoD_docs("α"))

$zerophase_docs
"""
function kaiser(n::Integer, α::Real; padding::Integer=0, zerophase::Bool=false)
    pf = 1.0/besseli(0,pi*α)
    makewindow(n, padding, zerophase) do x
        pf*besseli(0, pi*α*(sqrt(1 - (2x)^2)))
    end
end

# Discrete prolate spheroid sequences (Slepian tapers)
#
# See Gruenbacher, D. M., & Hummels, D. R. (1994). A simple algorithm
# for generating discrete prolate spheroidal sequences. IEEE
# Transactions on Signal Processing, 42(11), 3276-3278.
"""
$dpss_winplot

    dpss(n::Integer, nw::Real, ntapers::Integer=iceil(2*nw)-1;
         padding::Integer=0, zerophase::Bool=false)

The first `ntapers` discrete prolate spheroid sequences (Slepian
tapers) as an `n` × `ntapers` matrix. The signs of the tapers
follow the convention that the first element of the skew-symmetric
(odd) tapers is positive. The time-bandwidth product is given by
`nw`.

The DPSS window maximizes the energy concentration in the main lobe.

$zerophase_docs
"""
function dpss(n::Integer, nw::Real, ntapers::Integer=ceil(Int, 2*nw)-1;
              padding::Integer=0, zerophase::Bool=false)
    if isodd(n) && zerophase
        throw(ArgumentError("`dpss` does not currently support odd-length zerophase windows"))
    end
    # TODO there's probably a more efficient way to compute the zero-phase
    # version of this
    if zerophase
        n += 1
    end
    0 < ntapers <= n || error("ntapers must be in interval (0, n]")
    0 <= nw < n/2 || error("nw must be in interval [0, n/2)")

    # Construct symmetric tridiagonal matrix
    v = cospi(2*nw/n)
    dv = Vector{Float64}(undef, n)
    ev = Vector{Float64}(undef, n - 1)
    @inbounds dv[1] = v * abs2((n - 1) / 2)
    @inbounds @simd for i = 1:(n-1)
        dv[i + 1] = v * abs2((n - 1) / 2 - i)
        ev[i] = 0.5 * (i * n - i^2)
    end
    mat = SymTridiagonal(dv, ev)

    # Get tapers
    eigvec = eigen!(mat, n-ntapers+1:n).vectors
    rv = reverse(eigvec, dims=2)::Matrix{Float64}

    # Slepian's convention; taper starts with a positive element
    sgn = ones(size(rv, 2))
    for i = 2:2:size(rv, 2)
        s = zero(eltype(rv))
        for j = 1:n
            s = sign(rv[j, i])
            s != 0 && break
        end
        @assert s != 0
        sgn[i] = s
    end
    rmul!(rv, Diagonal(sgn))

    if zerophase
        # remove the last element (should be equal to the first element)
        rv = rv[1:end-1, :]
    end

    if padding > 0
        rv = [rv; zeros(padding, ntapers)]
    end

    if zerophase
        # TODO: figure out how to handle odd-length zerophase windows. See the
        # tests for one approach, but upsampling by 2 doesn't work here becasue
        # the overall amplitude varies with `n`
        rv = ifftshift(rv, 1)
    end

    rv
end

# Eigenvalues of DPSS, following Percival & Walden p. 390, exercise 8.1
# See also implementation in MNE:
# https://github.com/mne-tools/mne-python/blob/d7082cf909ccab581667bc1f1ed3c23e6a24b567/mne/time_frequency/multitaper.py#L226
"""
    dpsseig(A, nw)

Eigenvalues of the DPSS matrix, representing the ratios of the
power within the main lobe to the total power (main and sidelobes).
`A` is the output of [`dpss`](@ref), and `nw` is the
time-bandwidth product provided to [`dpss`](@ref) as input.
"""
function dpsseig(A::Matrix{Float64}, nw::Real)
    0 <= nw < size(A, 1)/2 || error("nw must be in interval [0, n/2)")

    w = nw/size(A, 1)

    # Compute coefficients
    seq = Vector{Float64}(undef, size(A, 1))
    seq[1] = 1.0
    for i = 1:size(A, 1)-1
        seq[i+1] = 2 * sinc(2w*i)
    end

    q = Vector{Float64}(undef, size(A, 2))
    nfft = nextfastfft(2*size(A, 1)-1)

    tmp1 = Vector{Float64}(undef, nfft)
    tmp2 = Vector{Complex{Float64}}(undef, nfft >> 1 + 1)
    p1 = plan_rfft(tmp1)
    p2 = plan_brfft(tmp2, nfft)

    for i = 1:size(A, 2)
        fill!(tmp1, 0)
        copyto!(tmp1, 1, A, (i-1)*size(A, 1)+1, size(A, 1))
        mul!(tmp2, p1, tmp1)
        for j = 1:length(tmp2)
            @inbounds tmp2[j] = abs2(tmp2[j])
        end
        mul!(tmp1, p2, tmp2)

        eig = 0.0
        for j = 1:size(A, 1)
            eig += seq[j]*tmp1[j]
        end
        q[i] = 2w * eig / nfft
    end

    q
end

# tensor product window functions in 2-d, defined as w(x,y) = w(x) \otimes w(y)

# convert the given argument to a 2-Tuple so we can distribute it to the
# input windows when making a 2D window
argdup(arg::Tuple) = arg
argdup(arg::Real) = (arg, arg)
const IntegerOr2 = Union{Tuple{<:Integer, <:Integer}, Integer}
const RealOr2 = Union{Tuple{<:Real, <:Real}, Real}
const BoolOr2 = Union{Tuple{Bool, Bool}, Bool}

for func in (:rect, :hanning, :hamming, :cosine, :lanczos,
             :triang, :bartlett, :bartlett_hann, :blackman)
    @eval begin
        function $func(dims::Tuple; padding::IntegerOr2=0,
                                    zerophase::BoolOr2=false)
            length(dims) == 2 || throw(ArgumentError("`dims` must be length 2"))
            paddings = argdup(padding)
            zerophases = argdup(zerophase)
            w1 = $func(dims[1], padding=paddings[1], zerophase=zerophases[1])
            w2 = $func(dims[2], padding=paddings[2], zerophase=zerophases[2])
            w1 * w2'
        end
    end
end

for func in (:tukey, :gaussian, :kaiser)
    @eval begin
        function $func(dims::Tuple, arg::RealOr2;
                       padding::IntegerOr2=0, zerophase::BoolOr2=false)
            length(dims) == 2 || throw(ArgumentError("`dims` must be length 2"))
            args = argdup(arg)
            paddings = argdup(padding)
            zerophases = argdup(zerophase)
            w1 = $func(dims[1], args[1]; padding=paddings[1], zerophase=zerophases[1])
            w2 = $func(dims[2], args[2]; padding=paddings[2], zerophase=zerophases[2])
            w1 * w2'
        end
    end
end

end # end module definition