# Code derived from metabolicEP (https://github.com/anna-pa-m/Metabolic-EP)

function compute_mom5d(xinf, xsup)

    minval = min(abs(xinf), abs(xsup))
    sgn = sign(xinf*xsup)

    if xsup - xinf < 1e-8
        return (0.5*(xsup + xinf),-1.)
    end

    if minval <= 6. || sgn <= 0
        ϕsup   = ϕ(xsup)
        Φsup   = Φ(xsup)
        ϕinf   = ϕ(xinf)
        Φinf   = Φ(xinf)
        scra1 = (ϕinf - ϕsup)/(Φsup - Φinf)
        scra2 = (xinf * ϕinf - xsup*ϕsup)/(Φsup -Φinf)
        scra12 = scra2 - scra1^2
        return scra1, scra12
    else
        delta2 = (xsup^2 - xinf^2)*0.5
        if delta2 > 40.
            scra1 = xinf^5/(3 - xinf^2 + xinf^4)
            scra2 = xinf^6/(3 - xinf^2 + xinf^4)
        else
            scra1 = (xinf*xsup)^5 * (1. - exp(delta2)) / (-exp(delta2)*(3.0-xinf^2 + xinf^4)*xsup^5 + xinf^5*(3-xsup^2 + xsup^4))
            scra2 = (xinf*xsup)^5 * (xsup - xinf*exp(delta2)) / (-exp(delta2)*(3.0-xinf^2 + xinf^4)*xsup^5 + xinf^5*(3-xsup^2 + xsup^4))
        end
        scra12 = scra2 - scra1^2
        isnan(scra1) || isnan(scra2) || isnan(scra12) && println("scra1 = $scra1 scra2 = $scra2")
        !isfinite(scra1) ||  !isfinite(scra2) && println("scra1 = $scra1 scra2 = $scra2")
        return scra1, scra12
    end
end


module Interpolations

export linerinterp

function linerinterp(grid,vals)

    function funfun(x)

        if x < grid[1]

            return (vals[2]-vals[1])/(grid[2]-grid[1])*(x-grid[1])+vals[1]

        end

          if  grid[end] <= x

               return (vals[end]-vals[end-1])/(grid[end]-grid[end-1])*(x-grid[end-1])+vals[end-1]

        end

        if grid[1] <= x < grid[end]

            i = searchsortedlast(grid,x)

            y = (vals[i+1]-vals[i])/(grid[i+1]-grid[i])*(x-grid[i])+vals[i]

            return y

        end

    end

    return funfun

end

end


abstract type ArchimedeanCopula <: AbstractCopula end

function rosenblatt(M::AbstractMatrix, c::ArchimedeanCopula)
    u1 = M[:, 1]
    u2 = M[:, 2]

    return [u1 (φ².(φ⁻¹.(u1, c) + φ⁻¹.(u2, c), c) ./ φ².(φ⁻¹.(u1, c), c))]
end

function cdf(c::ArchimedeanCopula, u1::Real, u2::Real)
    return φ(φ⁻¹(u1, c) + φ⁻¹(u2, c), c)
end

function density(c::ArchimedeanCopula, u1::Real, u2::Real)
    return φ²(φ⁻¹(u1, c) + φ⁻¹(u2, c), c) * ∇φ⁻¹(u1, c) * ∇φ⁻¹(u2, c)
end

function sample(c::ArchimedeanCopula, n::Int64)
    @assert n > 0
    return inverse_rosenblatt(rand(n, 2), c)
end

function (c::ArchimedeanCopula)(
    X::ContinuousUnivariateDistribution, Y::ContinuousUnivariateDistribution
)
    # Construct a joint distribution using the copula
    return Joint(X, Y, c)
end

include("independence.jl")
include("clayton.jl")


#-------------split-op.jl------------------------------------------------------#
#
# Purpose: A simple split-operator method for solving the schrodinger equation
#          The split operator method involves continuous transformations from
#              real space to phase space, before each tranformation, we multiply
#              by operators that exist either completely in real space or in
#              phase space. This allows us to solve Psi = Psi_0*e^(iHt)
#
#------------------------------------------------------------------------------#

using PyPlot

# Module to hold all parameters for simulation
module Par
    const xmax = 10.0
    const res = 512
    const dt = 0.05
    const timesteps = 1000
    const dx = 2*xmax / res
    const x = -xmax+dx:dx:xmax
    const dk = pi / xmax
    const k = vcat(0:res/2-1, -res/2:-1)*dk
end

# Type to hold all operators
type Operators
    V::Vector{Complex{Float64}}
    PE::Vector{Complex{Float64}}
    KE::Vector{Complex{Float64}}
    wfc::Vector{Complex{Float64}}
end

# Function to initialize the wfc and potential
function init(voffset::Float64, wfcoffset::Float64)
    V = 0.5 * (Par.x - voffset).^2
    wfc = 3* exp(-(Par.x-wfcoffset).^2/2)
    PE = exp(-0.5*im*V*Par.dt)
    KE = exp(-0.5*im*Par.k.^2*Par.dt)

    opr = Operators(V, PE, KE, wfc)
end

# Function for the split-operator loop
function split_op(opr::Operators)

    for i = 1:Par.timesteps
        # Half-step in real space
        opr.wfc = opr.wfc.*opr.PE

        # fft to phase space
        opr.wfc = fft(opr.wfc)

        # Full step in phase space
        opr.wfc = opr.wfc.*opr.KE

        # ifft back
        opr.wfc = ifft(opr.wfc)

        # final half-step in real space
        opr.wfc = opr.wfc.*opr.PE

        # plotting density and potential
        density = abs2(opr.wfc)

        clf()
        plot(density)
        plot(opr.V)
        sleep(0.01)
    end
end

# main function
function main()
    opr = init(2.0,-1.0)
    split_op(opr)
end

main()


println("# HDP TESTS...")

# run demo code at first
println(" * try demo...")

# define size of bar image
img_size = 5

# define amount of noise
noise_level = 0.01

# define probabilities of generating a particular number of bars
num_per_mixture = [0 ones(1,3)] ./ 3

# define number of groups (J)
num_group = 40

# define number of data items drawn from each group
num_data = 50;

# generate data
(X, bars) = generateBarsDataset(img_size, noise_level, num_per_mixture, num_group, num_data);

# Dimensionality of the data
D = img_size * img_size

# We assume a Multinomial Distribution with a Dirichlet Prior as base distribution
H = MultinomialDirichlet(D, 1.0)

# train a Hierarchical Dirichlet Process Mixture Model
# guessing 10 shared Distributions
models = train(HDP(H), Gibbs(), RandomInitialisation(k = 10), X);

# access the log likelihood values of all iteration
LLH = map(model -> model.energy, models)

# access the weights of all iterations
W = map(model -> model.weights, models)

# access the distributions of all iterations
D = map(model -> model.distributions, models)

println(" * test randnumtable()")
weights = ones(10, 10)
table = ones(Int, 10, 10)

@test weights == BNP.randnumtable(weights, table)

println("# FINISHED")


# Adapted from:
# SOSDEMO10 --- Set containment
# Section 3.10 of SOSTOOLS User's Manual

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

    eps = 1e-6
    p = x[1]^2+x[2]^2
    gamma = 1
    g0 = 2*x[1]
    theta = 1

    m = SOSModel(factory)

    # FIXME s should be sos ?
    # in SOSTools doc it is said to be SOS
    # but in the demo it is not constrained so
    Z = monomials(x, 0:4)
    @variable m s Poly(Z)

    Z = monomials(x, 2:3)
    @variable m g1 Poly(Z)

    Sc = [theta^2-s*(gamma-p) g0+g1; g0+g1 1]

    @SDconstraint m Matrix(eps * I, 2, 2) ⪯ Sc

    JuMP.optimize!(m)

    # Program is feasible, { x |((g0+g1) + theta)(theta - (g0+g1)) >=0 } contains { x | p <= gamma }
    @test JuMP.primal_status(m) == MOI.FeasiblePoint || JuMP.primal_status(m) == MOI.NearlyFeasiblePoint

    #@show JuMP.value(s)
    #@show JuMP.value(g1)
end


module SymbolicWedderburn

using LinearAlgebra
using SparseArrays
using Primes

using Cyclotomics
using GroupsCore
using PermutationGroups
using StarAlgebras

export symmetry_adapted_basis, WedderburnDecomposition
export basis,
    # degree, # too common name to export
    direct_summands,
    invariant_vectors,
    issimple,
    multiplicity

include("Characters/Characters.jl")
using .Characters
import .Characters: row_echelon_form!
import .Characters.FiniteFields

include("util.jl")
include("actions.jl")
include("action_characters.jl")
include("matrix_projections.jl")
include("minimal_projections.jl")
include("direct_summands.jl")
include("sa_basis.jl")
include("wedderburn_decomposition.jl")

end # module


function nlfeast_opt!(T, X::AbstractMatrix{ComplexF64}, nodes::Integer, iter::Integer;
    c=complex(0.0,0.0), r=1.0, debug=false, ϵ=0.05)

    N, m₀ = size(X)
    IM = Matrix(I, m₀, m₀)
    Λ = zeros(ComplexF64, m₀)
    θ = LinRange(π/nodes, 2*π-π/nodes, nodes)
    Q₀, Q₁ = zeros(ComplexF64, N, m₀), zeros(ComplexF64, N, m₀)
    Tinv, R = similar(X, ComplexF64), similar(X, ComplexF64)
    resolvent = similar(Λ)

    A, B = zeros(ComplexF64, m₀, m₀), zeros(ComplexF64, m₀, m₀)

    alpha, beta = similar(Λ), similar(Λ)
    vl, vr = zeros(ComplexF64, 1, m₀), zeros(ComplexF64, m₀, m₀)
    work, rwork = zeros(ComplexF64, 1), zeros(Float64, 8*m₀)

    TLU = [qr(T(r*exp(ang*im)+c)) for ang in θ]

    for i=1:nodes
	z = (r*exp(θ[i]*im)+c)
	Tinv .= (TLU[i] \ X) .* (r*exp(θ[i]*im)/nodes)
	Q₀ .+= Tinv
	Q₁ .+= Tinv .* z
    end

    mul!(A, X', Q₁)
    mul!(B, X', Q₀)
    lapack_ggev!(A, B, alpha, beta, vl, vr, work, rwork)
    mul!(X, Q₀, vr)
    Λ .= alpha ./ beta

    for nit=1:iter

        update_R!(X, R, Λ, T)

        Q₀ .= 0
        Q₁ .= 0

        for i=1:nodes
            z = (r*exp(θ[i]*im)+c)
	    resolvent .= (1 ./(z .- Λ))
	    Tinv .= X
	    Tinv .-= (TLU[i]\R)
	    rmul!(Tinv, (r*exp(θ[i]*im)/nodes) .* Diagonal(resolvent))
	    Q₀ .+= Tinv
	    Q₁ .+= Tinv .* z
        end

        mul!(A, X', Q₁)
        mul!(B, X', Q₀)
        lapack_ggev!(A, B, alpha, beta, vl, vr, work, rwork)
        mul!(X, Q₀, vr)
        Λ .= alpha ./ beta

        if debug println(nit) end
    end

    normalize!(X)
    Λ, X, residuals(R, Λ, T)
end


##=================####==============##
# Opinion dynamics
##=================####==============##

"""
    HK_od_model!(agent;kwargs...)
Hegselmann-Krause model
Updates `agent` cooperation opinion as a function of its neighbors
ϵ -> Bounded confidence
"""
function HK_od_model!(agents, ag_id)

    # Agent to be updated
    ag = agents[ag_id]

    # Agent's confidence set
    ag_conf_set = [
        agents[neigh]
        for neigh in ag.contacts_t if abs(ag.p_cop - agents[neigh].p_cop) <= ag.ϵ
    ]

    if !isempty(ag_conf_set)
        # Mean opinion of agent's confidence set
        conf_set_op = mean([neigh.p_cop for neigh in ag_conf_set])

        # Fraction of infected agents in the confidence set
        ϕ_I =
            count(x -> x == "I", [neigh.state for neigh in ag_conf_set]) / length(ag_conf_set)

        # New agents opninion
        ag.p_cop = conf_set_op + 0.1 * ϕ_I * (1.0 - conf_set_op)
        #  ag.p_cop = conf_set_op
    end
end

##=================####==============##


const logkappa_spl = opal_test_spl()

function drdm(m, r, l, P, T, rho, mu)
    #=
    Derivative of r with respect to m, from mass conservation.
    =#
    return 1 / (4 * pi * r .^ 2 .* rho)
end # drdm

function dPdm(m, r, l, P, T)
    #=
    Derivative of P with respect to m, from hydrostatic equilibrium.
    =#
    return -G * m / (4 * pi * r .^ 4)
end # dPdm

function dldm(m, r, l, P, T, rho, X, Y)
    #=
    Derivative of l with respect to m, from energy conservation.
    =#
    return epsilon_pp(rho, T, X) + epsilon_cno(rho, T, X, Y)
end # dldm

function dTdm(m, r, l, P, T, mu, kappa)
    #=
    Derivative of T with respect to m, from... thermodynamics.
    =#
    grad_rad = ∇_rad(m, l, P, T, kappa)
    grad_ad = ∇_ad()
    stable =  grad_rad < grad_ad
    if stable
        # @debug("dTdm: stable to convection")
        grad = grad_rad
    else
        # @debug("dTdm: unstable to convection")
        grad = grad_ad
    end # if
    factor = - G .* m .* T ./ (4 * pi .* r .^ 4 .* P)
     return factor .* grad
end # dTdm


mutable struct VerticalDiffusion

    amo :: AdvancedMatrixOperators

    eT_length    :: Int64

    eT_send_T    :: AbstractArray{Float64, 2}
    T_send_eT    :: AbstractArray{Float64, 2}
    eT_I_eT      :: AbstractArray{Float64, 2}

    K_iso :: Float64
    K_cva :: Float64

    K_func :: Function

    wksp

    function VerticalDiffusion(
        amo   :: AdvancedMatrixOperators;
        K_iso :: Float64,
        K_cva :: Float64,
    )

        cvtDiagOp = (a,) -> spdiagm(0 => view(a, :))
        Nz, Ny, Nx = amo.bmo.T_dim

        # Create coversion matrix and its inverse
        T_num        = reshape(amo.T_mask_T * collect(1:amo.bmo.T_pts), Nz, Ny, Nx)
        active_T_num = T_num[ T_num .!= 0.0 ]
        eT_send_T    = amo.bmo.T_I_T[active_T_num, :]
        T_send_eT    = eT_send_T' |> sparse
        eT_I_eT      = eT_send_T * T_send_eT

        eT_length = size(eT_I_eT)[1]

        #K = spdiagm(0 => ones(Float64, amo.bmo.W_pts))

        K_func = (dbdz,) -> (dbdz > 0.0) ? K_iso : K_cva

        # iEBD = inverse of Euler Backward Diffusion ( I - Δt * ∇K∇ )
        wksp = (
            above_W    = zeros(Float64, amo.bmo.W_pts),
            below_W    = zeros(Float64, amo.bmo.W_pts),
            Δ_W        = zeros(Float64, amo.bmo.W_pts),
            K_W        = zeros(Float64, amo.bmo.W_pts),
            W_K_W      = spdiagm(0 => ones(Float64, amo.bmo.W_pts)),
            W_K∇_T     = spzeros(Float64, amo.bmo.W_pts, amo.bmo.T_pts),
            T_∇K∇_T    = spzeros(Float64, amo.bmo.T_pts, amo.bmo.T_pts),
            eT_∇K∇_T   = spzeros(Float64, eT_length, amo.bmo.T_pts),
            eT_∇K∇_eT  = spzeros(Float64, eT_length, eT_length),
            eT_iEBD_eT = spzeros(Float64, eT_length, eT_length),
            rhs_eT     = zeros(Float64, eT_length),
            lhs_eT     = zeros(Float64, eT_length),
        )

        return new(
            amo,
            eT_length,
            eT_send_T,
            T_send_eT,
            eT_I_eT,
            K_iso,
            K_cva,
            K_func,
            wksp,
        )

    end
end

function calOp_vdiff(
    vd   :: VerticalDiffusion,
    b    :: AbstractArray, # of grid :T
    HMXL :: AbstractArray, # of grid :sT
)

    amo = vd.amo
    bmo = amo.bmo
    gd  = amo.gd

    dbdz = vd.amo.W_∂z_T * view(b, :)
    HMXL = reshape(HMXL, 1, gd.Nx, gd.Ny)

    K_W = reshape( vd.K_func.(dbdz), bmo.W_dim...)
    K_W[gd.z_W .> - HMXL] .= vd.K_cva

    # below this W-pts the temperature is below freezing point.
    # reasoning please see stepColumn discussing about Q_FRZHEAT / frzmlt
    #below_frz_below = (bmo.W_UP_T * view(TEMP, :)) .< T_sw_frz
    #K_W[below_frz_below] .= vd.K_cva

    #K_W .= 0
    op_vdiff = sparse(vd.amo.T_DIVz_W * vd.amo.W_mask_W * spdiagm( 0 => view(K_W,:)) * vd.amo.W_∂z_T)

    dropzeros!(op_vdiff)

    return op_vdiff
end


function MPT_SVD(arrayaux::Array{Float64},middlesizes::Matrix{Int64},acc::Float64)
    tensor = Array{Float64};
    array  = copy(arrayaux);
    if size(middlesizes,1) == 1
        tensor = reshape(array,Tuple(middlesizes));
        return TT_SVD(tensor,acc);
    else
        sizes   = Tuple(reshape(middlesizes',(length(middlesizes),1)));
        array   = reshape(array,sizes);
        permind = Tuple([ (i-1)*size(middlesizes,2)+j for j in 1:size(middlesizes,2) for i in 1:size(middlesizes,1) ]);
        array   = permutedims(array,permind);
        resind  = Tuple([prod(col) for col in eachcol(middlesizes)]);
        tensor  = reshape(array,resind)
        mps,err = TT_SVD(tensor,acc);
        rnks    = rank(mps);
        return MPT( [reshape(mps[i],(rnks[i][1], middlesizes[:,i]..., rnks[i][2])) for i = 1:order(mps)] ), err
    end
end

# TT-ALS for diagonal Matrix
function MPT_SVD(diagonal::Diagonal,middlesizes::Matrix,acc::Float64)
    # initial ttm with diagonal core
    N      = length(middlesizes);
    tensor = reshape(diag(diagonal),Tuple(middlesizes));
    tt,err = TT_SVD(tensor,acc);
    cores  = Vector{Array{Float64,4}}(undef,N);
    rnks   = rank(tt);
    for i = 1:N
        cores[i] = zeros(rnks[i][1],middlesizes[i],middlesizes[i],rnks[i][2])
        for ri = 1:rnks[i][1]
            for rii = 1:rnks[i][2]
                cores[i][ri,:,:,rii] = Matrix(Diagonal(tt[i][ri,:,rii]))
            end
        end
    end
    return MPT(cores)
end

# TT-ALS for symmetric Matrix
function MPT_SVD(sym::Symmetric,middlesizes::Matrix,acc::Float64)
    # initial ttm with diagonal core
    d          = Int(length(middlesizes)/2)
    symM       = Matrix(sym)
    tensor     = reshape(symM,Tuple(middlesizes));
    Uttm,S,err = TT_SVD(tensor,d,acc);
    r          = rank(Uttm)

    cores   = Vector{Array{Float64,3}}(undef,2*d)
    coresL  = Vector{Array{Float64,4}}(undef,d)
    coresU  = Vector{Array{Float64,4}}(undef,d)

    di     = collect(d:-1:1)
    for i = 1:d
        cores[i]   = Uttm[i]
        cores[d+i] = permutedims(Uttm[di[i]],[3,2,1])

        if i<d
            coresL[i]  = reshape(Uttm[i],(size(Uttm[i],1),size(Uttm[i],2),1,size(Uttm[i],3)))
            coresU[i]  = coresL[i]
        end
    end

    coresU[d]  = reshape(cores[d],(size(cores[d],1),size(cores[d],2),size(cores[d],3),1))
    cores[d]   = nmodeproduct(Matrix(sqrt.(Diagonal(S))),cores[d],3)
    cores[d+1] = nmodeproduct(Matrix(sqrt.(Diagonal(S))),cores[d+1],1)
    coresL[d]  = reshape(cores[d],(size(cores[d],1),size(cores[d],2),size(cores[d],3),1))

    return MPT(cores),MPT(coresL),MPT(coresU),S,err
end


let
    n = 50
    p = 3
    t = (1:n)/n
    Ut,Vt = egrss.generators(t,p);
    Wt,c = egrss.potrf(Ut,Vt,1e-2*ones(n))
    Yt,Zt = egrss.trtri(Ut,Wt,c)

    Lref = cholesky(egrss.full(Ut,Vt) + Diagonal(1e-2*ones(n))).L
    Lref_inv = LowerTriangular(tril(inv(Lref)));

    L_inv = egrss.full_tril(Yt,Zt,1.0./c)
    @test ( isapprox(L_inv,Lref_inv) )

    L_inv = egrss.trtri2(Ut,Wt,c)
    @test ( isapprox(L_inv,Lref_inv) )
end


module Unfolding

using Clustering: dbscan
using DelimitedFiles
using Distances
using ImageMorphology: thinning
using Graphs: connected_components, dijkstra_shortest_paths
using LinearAlgebra: Diagonal, Symmetric, cross, dot, eigen
using MultivariateStats: PCA, dmat2gram, fit, projection
using NearestNeighbors
using Optim
using Random
using Setfield
using SimpleWeightedGraphs
using StatsBase: Weights, mean, mean!, quantile, sample
using WriteVTK

include("coords_optimization.jl")
include("data_handling.jl")
include("landmark_isomap.jl")
include("reference_points.jl")
include("unfold.jl")

export
    coords,
    errors,
    getreference,
    landmark_isomap,
    to_vtk,
    unfold
end


p = ℝ³(rand(3))
q = Quaternion(rand() * 2pi, normalize(ℝ³(rand(3))))
q₀ = Quaternion(0, normalize(ℝ³(rand(3))))
q2 = normalize(Quaternion(rand(4)))
points = [ℝ³(rand(3)) for i in 1:5]
i = ℝ³(1, 0, 0)
j = ℝ³(0, 1, 0)
k = ℝ³(0, 0, 1)
g = getrotation(i, j)

@test typeof(rotate(p, q)) == typeof(p)
@test isapprox(rotate(p, q₀), p) # rotations with angle zero
@test typeof(rotate(points, q)) == typeof(points) # Array input
@test isapprox(g, Quaternion(pi / 2, k))

q = Biquaternion(Quaternion(rand() * 2pi - pi, ℝ³(rand(3))), ℝ³(rand(3)))
transformed = applyconfig(points, q)

@test typeof(transformed) == typeof(points)
@test size(transformed) == size(points)
@test isapprox(transformed, points) == false


summation(num) = (num*(num+1)) ÷ 2


function numerical_integration(x, y)
    sum(diff(x) .* midpoints(y))
end
function cdf_distance(F_X, F_Y, x, p)
    if p == Inf
        return maximum(abs.(F_X .- F_Y))
    end
    diff_p = abs.(F_X .- F_Y).^p
    return numerical_integration(x, diff_p)^(1/p)
end
function cdf_distance(Xdistr::UnivariateDistribution, Ydistr::UnivariateDistribution, x, p)
    F_X = cdf.(Xdistr, x)
    F_Y = cdf.(Ydistr, x)
    return cdf_distance(F_X, F_Y, x, p)
end

function distance_test_statistic(p, method::DeconvolutionMethod, xgrid::AbstractVector)
    function ϕ(X, Y, σ_X, σ_Y)
        ϵ_X = Normal.(0.0, σ_X)
        Xdistr_hat = decon(method, X, ϵ_X)

        ϵ_Y = Normal.(0.0, σ_Y)
        Ydistr_hat = decon(method, Y, ϵ_Y)
        return cdf_distance(Xdistr_hat, Ydistr_hat, xgrid, p)
    end
    return ϕ
end

KS_test_statistic(X, Y, σ_X, σ_Y) = HypothesisTests.ApproximateTwoSampleKSTest(X, Y).δ


module Quantum

using LinearAlgebra
import Base: ==, length

export Swap, X, Y, Z, H, S, Sdag, T, Tdag, Rx, Ry, Rz, R, CNOT, CZ, CR, CRk
export QuantumRegister

mutable struct QuantumRegister
    qubit_product::Vector{Complex{Float64}}
    QuantumRegister(length::Int) = new(vcat([1], zeros(Complex{Float64}, 2^length - 1)))
end

==(register::QuantumRegister, vector::Vector) = register.qubit_product == vector
==(vector::Vector, register::QuantumRegister) = vector == register.qubit_product
==(a::QuantumRegister, b::QuantumRegister) = a.qubit_product == b.qubit_product

# Because `QuantumRegister` saves the tensor product of all its qubits,
# its length is 2^n (where n is the amount of qubits) and needs to be modified
length(register::QuantumRegister) = log2size(register.qubit_product)

# These functions are defined here because they give an `UndefVarError` bellow

single_qubit_gate(matrix::Matrix) =
    (register::QuantumRegister, at::Int) -> apply!(register, matrix, at)

controlled_gate(matrix::Matrix) = function(register::QuantumRegister, control::Int, target::Int)
    _control = control
    _target = target

    if(control == length(register))
        _control -= 1
        Swap(register, control, _control)
        if(_control == target)
            _target = control
        end
    end

    if(_control + 1 != _target)
        _target = _control + 1
        Swap(register, target, _target)
    end

    apply!(register, controlled_matrix(matrix), _control)

    Swap(register, _target, target)

    if(_control == _target)
        Swap(register, _control, control)
    end

    register
end

# This is where the real gates start

Swap(register::QuantumRegister, from::Int, to::Int) =
    apply!(register, swap_matrix(abs(to - from) + 1), min(from, to))

X = single_qubit_gate([0 1 ; 1 0])

Y = single_qubit_gate([0 -im ; im 0])

Z = single_qubit_gate([1 0 ; 0 -1])

H = single_qubit_gate([1/√2 1/√2 ; 1/√2 -1/√2])

S = single_qubit_gate([1 0 ; 0 im])

Sdag = single_qubit_gate([1 0 ; 0 -im])

T = single_qubit_gate([1 0 ; 0 exp(im * pi / 4)])

Tdag = single_qubit_gate([1 0 ; 0 exp(-im * pi / 4)])

Rx(register::QuantumRegister, at::Int, θ::Float64) =
    single_qubit_gate([cos(θ) -im * sin(θ) ; -im * sin(θ) cos(θ)])(register, at)

Ry(register::QuantumRegister, at::Int, θ::Float64) =
    single_qubit_gate([cos(θ / 2) -im * -sin(θ / 2) ; -im * sin(θ) cos(θ / 2)])(register, at)

Rz(register::QuantumRegister, at::Int, θ::Float64) =
    single_qubit_gate([0 1 ; 0 exp(im * θ)])(register, at)

# R is often used to denote Rz
R = Rz

CNOT = controlled_gate([0 1 ; 1 0])

CZ = controlled_gate([1 0 ; 0 -1])

CR(register::QuantumRegister, control::Int, target::Int, θ::Float64) =
    controlled_gate([0 1 ; 0 exp(im * θ)])(register, control, target)

CRk(register::QuantumRegister, control::Int, target::Int, k::Float64) =
    CR(register, control, target, 2π / 2 ^ k)

# Helper functions for the gates

controlled_matrix(matrix::Matrix) = kron([1 0 ; 0 0], I₂) + kron([0 0 ; 0 1], matrix)

function apply!(register::QuantumRegister, matrix::Matrix, from::Int)
    register.qubit_product = pad_matrix(matrix, length(register), from) * register.qubit_product

    register
end

pad_matrix(matrix::Matrix, tosize::Int, from::Int) =
    # kron together identity matrices with `matrix` in index `from`
    foldl(kron, vcat(fill(I₂, from - 1), [matrix], fill(I₂, tosize - from - log2size(matrix) + 1)))

function swap_matrix(span::Int)
    matrix = I + zeros(2^span, 2^span)

    for i in 1:2^(span - 1)
        if iseven(i)
            swap_rows!(matrix, i, i + 2^(span - 1) - 1)
        end
    end

    matrix
end

function swap_rows!(matrix::Matrix, x::Int, y::Int)
    for i in 1:size(matrix, 2)
        temp = matrix[x, i]
        matrix[x, i] = matrix[y, i]
        matrix[y, i] = temp
    end

    matrix
end

log2size(array::Array) = Int(log2(size(array, 1)))

I₂ = [1 0 ; 0 1]

end # module


import Gen
import Distributions
import LinearAlgebra#: norm, cross, det
import Rotations
import StaticArrays

const Rot2 = Rotations.RotMatrix{2,Float64,4}

function Base.rand(::Type{Rot2})
    a = Gen.uniform(-pi, pi)
    return Rotations.RotMatrix{2,Float64}(a)
end

total_area(::Type{Rot2}) = 2 * pi

@inline function _to_array(x::Rot2)
    return Float64[x[1,1], x[2,1]]
end

export Rot2

################################################
# uniform distribution (Haar measure on SO(2)) #
################################################

struct UniformRot2 <: Gen.Distribution{Rot2} end

Gen.logpdf(::UniformRot2, x::Rot2) = -log(total_area(Rot2))

function Gen.logpdf_grad(::UniformRot2, x::Rot2)
    error("Not implemented") # TODO implement
    return (nothing,)
end

Gen.random(::UniformRot2) = rand(Rot2)

Gen.has_output_grad(::UniformRot2) = false

Gen.has_argument_grads(::UniformRot2) = ()

(::UniformRot2)() = Gen.random(UniformRot2())

const uniform_rot2 = UniformRot2()

export uniform_rot2

################################################
# von Mises Fisher based distribution on SO(2) #
################################################

struct VMFRot2 <: Gen.Distribution{Rot2} end

const von_mises_rot2 = VMFRot2()

function Gen.logpdf(::VMFRot2, x::Rot2, mu::Rot2, k::Real)
    dist = Distributions.VonMisesFisher(_to_array(mu), k)
    x_unit_vec = _to_array(x)
    return Distributions.logpdf(dist, x_unit_vec)
end

function Gen.logpdf_grad(::VMFRot2, x::Rot2, mu::Rot2, k::Real)
    error("Not implemented") # TODO implement
    return (nothing,nothing,nothing)
end

function Gen.random(::VMFRot2, mu::Rot2, k::Real)
    dist = Distributions.VonMisesFisher(_to_array(mu), k)
    v = rand(dist)
    mat = StaticArrays.SMatrix{2,2,Float64,4}(v[1], v[2], -v[2], v[1])
    return Rotations.RotMatrix{2,Float64,4}(mat)
end

Gen.has_output_grad(::VMFRot2) = false

Gen.has_argument_grads(::VMFRot2) = (false, false)

(::VMFRot2)(mu, k) = Gen.random(VMFRot2(), mu, k)

export von_mises_rot2


## Motion Planning Benchmark
## =========================

# Construct a path from the start region to the target region
# which does not pass trhoug nay of the obstacles

typealias Point Lifted{Vector{Float64}}
typealias Vec Lifted{Vector{Float64}}
typealias Mat Lifted{Matrix{Float64}}
typealias PointPair Lifted{Matrix{Float64}}
typealias Edge Lifted{Matrix{Float64}}  #Parametric form of line is p1 + (p2 - p1)
typealias Scalar Lifted{Float64}

points_to_vec(p1::Point, p2::Point) = p1 - p2
points_to_vec(edge::PointPair) = points_to_vec(edge[:,1], edge[:,2])

# Parametric form of line is p1 + (p2 - p1)
points_to_parametric(p1::Point,p2::Point) = liftedarray([p1 points_to_vec(p2,p1)])
points_to_parametric(edge::PointPair) = points_to_parametric(edge[:,1], edge[:,2])
parametric_to_point(p::Edge, s::Scalar) = p[:,1] + s * p[:,2]

# Is the point in the box?
function ispointinpoly(point, box)
  r = (point[1] >= box[1,1]) &
    (point[1] <= box[2,1]) &
    (point[2] >= box[1,2]) &
    (point[2] <= box[2,2])
  r
end

# Linear Check
function avoid_obstacles(points, obs)
  observations = true
  for i = 1:size(points,2),ob in obs
    observations &= !point_in_poly(points[:,i],ob)
#     @show i, observations
  end
  observations
end

# Where if anywhere, along p does it intersect segment
function intersect_segments(ppos::Point, pdir::Vec, qpos::Point, qdir::Vec)
  w = ppos - qpos
  u = pdir
  v = qdir
  (v[2] * w[1] - v[1] * w[2]) / (v[1] * u[2] - v[2] * u[1])
end

function short_lines(points, length::Float64)
  true
end

function isvalid_path(points, obstacles, origin, dest)
  point_in_poly(points[:,1],origin) &
  ispointinpoly(points[:,end], dest) &
  avoid_obstacles(points,obstacles) &
  short_lines(points, .1)
end

function linearplan(points)
  origin = [0.0 0.0
           0.2 0.2]
  dest = [.9 .9
          1.0 1.0]

  box1 = [.3  .1
          .65 .7]
  box2 = [.05 .7
          .25 .99]

  obstacles = Matrix[box1,box2]
  good_path = isvalid_path(points,obstacles,origin,dest)
  points,good_path
end

## Nonlinear
## =========
function isvalid_pathnl(points, obstacles, origin, dest)
  ispointinpoly(points[:,1],origin) &
  ispointinpoly(points[:,end], dest) &
  avoid_obstacles_nl(points,obstacles) &
  short_lines(points, .1)
end

# Nonlinear Check
function avoid_obstacles_nl(points, obs)
  observations = true
  for i = 1:size(points,2)-1, ob in obs
    p = points[:,i]
    pdir = points[:,i+1] - points[:,i]

    obpos = ob[:,1]
    obdir = ob[:,2]
    s = intersect_segments(p, pdir, obpos, obdir)
    observations &= (s < 0) | (s > 1)
  end
  observations
end

function nonlinearplan(points)
  origin = [0.0 0.0
           0.2 0.2]
  dest = [.9 .9
          1.0 1.0]
  obstacles = Array[[8.01 3.01; 1.02 9],
                    [0.5 3.08; 2.02 9.04],
                    [0.0 9.99; 3 5.04]]
  obs = map(points_to_parametric, obstacles)
  good_path = isvalid_pathnl(points,obs,origin,dest)
  points,good_path
end

points = mvuniform(0.0, 10.0, 2,6)
points, condition = nonlinearplan(points)
@show condition.smt
# s = cond_sample_tlmh(points,condition,1)

# # ## Draw
# obstacles = Array[[8.01 3.01; 1.02 9],
#                   [0.5 3.08; 2.02 9.04],
#                   [0.0 9.99; 3 5.04]]
# a = make_compose_lines(obstacles)
# points = s[1]
# b = [Compose.line([pair(points[:,i]),pair(points[:,i+1])]) for i = 1:(size(points,2)-1)]
# b
# draw_lines(a,b)


#################### CSR

function _mul(A::SparseCSR{Tv,Ti}, x::Vector{Tv})::Vector{Tv} where {Tv, Ti}
    m, n = size(A)
    y = zeros(m)
    for i = 1:m
        for z = A.rowptr[i]:A.rowptr[i+1]-1
            j = A.colind[z]
            y[i] += A.val[z] * x[j]
        end
    end
    return y
end

function _mul(At::Adjoint{SparseCSR{Tv,Ti}}, x::Vector{Tv})::Vector{Tv} where {Tv, Ti}
    A = At.parent
    n, m = size(A)
    y = zeros(m)
    for j = 1:n
        for z = A.rowptr[j]:A.rowptr[j+1]-1
            i = A.colind[z]
            y[i] += A.val[z] * x[j]
        end
    end
    return y
end

###

function _mul(A::SparseCSR{Tv,Ti}, B::Matrix{Tv})::Matrix{Tv} where {Tv, Ti}
    m, k = size(A)
    k, n = size(B)
    C = zeros(m,n)
    for i = 1:m
        for z = A.rowptr[i]:A.rowptr[i+1]-1
            l = A.colind[z]
            for j = 1:n
                C[i,j] += A.val[z] * B[l,j]
            end
        end
    end
    return C
end

function _mul(A::SparseCSR{Tv,Ti}, Bt::LinearAlgebra.Adjoint{Tv,Matrix{Tv}})::Matrix{Tv} where {Tv, Ti}
    B = Bt.parent
    m, k = size(A)
    n, k = size(B)
    C = zeros(m,n)
    for i = 1:m
        for z = A.rowptr[i]:A.rowptr[i+1]-1
            l = A.colind[z]
            for j = 1:n
                C[i,j] += A.val[z] * B[j,l]
            end
        end
    end
    return C
end

function _mul(At::Adjoint{SparseCSR{Tv,Ti}}, B::Matrix{Tv})::Matrix{Tv} where {Tv, Ti}
    A = At.parent
    k, m = size(A)
    k, n = size(B)
    C = zeros(m,n)
    for l = 1:k
        for z = A.rowptr[l]:A.rowptr[l+1]-1
            i = A.colind[z]
            for j = 1:n
                C[i,j] += A.val[z] * B[l,j]
            end
        end
    end
    return C
end

function _mul(At::Adjoint{SparseCSR{Tv,Ti}}, Bt::LinearAlgebra.Adjoint{Tv,Matrix{Tv}})::Matrix{Tv} where {Tv, Ti}
    A = At.parent
    B = Bt.parent
    k, m = size(A)
    n, k = size(B)
    C = zeros(m,n)
    for l = 1:k
        for z = A.rowptr[l]:A.rowptr[l+1]-1
            i = A.colind[z]
            for j = 1:n
                C[i,j] += A.val[z] * B[j,l]
            end
        end
    end
    return C
end

###

function _mul(B::Matrix{Tv}, A::SparseCSR{Tv,Ti})::Matrix{Tv} where {Tv, Ti}
    m, k = size(B)
    k, n = size(A)
    C = zeros(m,n)
    for l = 1:k
        for z = A.rowptr[l]:A.rowptr[l+1]-1
            j = A.colind[z]
            for i = 1:m
                C[i,j] += A.val[z] * B[i,l]
            end
        end
    end
    return C
end

function _mul(Bt::LinearAlgebra.Adjoint{Tv,Matrix{Tv}}, A::SparseCSR{Tv,Ti})::Matrix{Tv} where {Tv, Ti}
    B = Bt.parent
    k, m = size(B)
    k, n = size(A)
    C = zeros(m,n)
    for l = 1:k
        for z = A.rowptr[l]:A.rowptr[l+1]-1
            j = A.colind[z]
            for i = 1:m
                C[i,j] += A.val[z] * B[l,i]
            end
        end
    end
    return C
end

function _mul(B::Matrix{Tv}, At::Adjoint{SparseCSR{Tv,Ti}})::Matrix{Tv} where {Tv, Ti}
    A = At.parent
    m, k = size(B)
    n, k = size(A)
    C = zeros(m,n)
    for j = 1:n
        for z = A.rowptr[j]:A.rowptr[j+1]-1
            l = A.colind[z]
            for i = 1:m
                C[i,j] += A.val[z] * B[i,l]
            end
        end
    end
    return C
end

function _mul(Bt::LinearAlgebra.Adjoint{Tv,Matrix{Tv}}, At::Adjoint{SparseCSR{Tv,Ti}})::Matrix{Tv} where {Tv, Ti}
    B = Bt.parent
    A = At.parent
    k, m = size(B)
    n, k = size(A)
    C = zeros(m,n)
    for j = 1:n
        for z = A.rowptr[j]:A.rowptr[j+1]-1
            l = A.colind[z]
            for i = 1:m
                C[i,j] += A.val[z] * B[l,i]
            end
        end
    end
    return C
end

#################### CSC

function _mul(A::SparseCSC{Tv,Ti}, x::Vector{Tv})::Vector{Tv} where {Tv, Ti}
    m, n = size(A)
    y = zeros(m)
    for j = 1:n
        for z = A.colptr[j]:A.colptr[j+1]-1
            i = A.rowind[z]
            y[i] += A.val[z] * x[j]
        end
    end
    return y
end

function _mul(At::Adjoint{SparseCSC{Tv,Ti}}, x::Vector{Tv})::Vector{Tv} where {Tv, Ti}
    A = At.parent
    n, m = size(A)
    y = zeros(m)
    for i = 1:m
        for z = A.colptr[i]:A.colptr[i+1]-1
            j = A.rowind[z]
            y[i] += A.val[z] * x[j]
        end
    end
    return y
end

###

function _mul(A::SparseCSC{Tv,Ti}, B::Matrix{Tv})::Matrix{Tv} where {Tv, Ti}
    m, k = size(A)
    k, n = size(B)
    C = zeros(m,n)
    for l = 1:k
        for z = A.colptr[l]:A.colptr[l+1]-1
            i = A.rowind[z]
            for j = 1:n
                C[i,j] += A.val[z] * B[l,j]
            end
        end
    end
    return C
end

function _mul(A::SparseCSC{Tv,Ti}, Bt::LinearAlgebra.Adjoint{Tv,Matrix{Tv}})::Matrix{Tv} where {Tv, Ti}
    B = Bt.parent
    m, k = size(A)
    n, k = size(B)
    C = zeros(m,n)
    for l = 1:k
        for z = A.colptr[l]:A.colptr[l+1]-1
            i = A.rowind[z]
            for j = 1:n
                C[i,j] += A.val[z] * B[j,l]
            end
        end
    end
    return C
end

function _mul(At::Adjoint{SparseCSC{Tv,Ti}}, B::Matrix{Tv})::Matrix{Tv} where {Tv, Ti}
    A = At.parent
    k, m = size(A)
    k, n = size(B)
    C = zeros(m,n)
    for i = 1:m
        for z = A.colptr[i]:A.colptr[i+1]-1
            l = A.rowind[z]
            for j = 1:n
                C[i,j] += A.val[z] * B[l,j]
            end
        end
    end
    return C
end

function _mul(At::Adjoint{SparseCSC{Tv,Ti}}, Bt::LinearAlgebra.Adjoint{Tv,Matrix{Tv}})::Matrix{Tv} where {Tv, Ti}
    A = At.parent
    B = Bt.parent
    k, m = size(A)
    n, k = size(B)
    C = zeros(m,n)
    for i = 1:m
        for z = A.colptr[i]:A.colptr[i+1]-1
            l = A.rowind[z]
            for j = 1:n
                C[i,j] += A.val[z] * B[j,l]
            end
        end
    end
    return C
end

###

function _mul(B::Matrix{Tv}, A::SparseCSC{Tv,Ti})::Matrix{Tv} where {Tv, Ti}
    m, k = size(B)
    k, n = size(A)
    C = zeros(m,n)
    for j = 1:n
        for z = A.colptr[j]:A.colptr[j+1]-1
            l = A.rowind[z]
            for i = 1:m
                C[i,j] += A.val[z] * B[i,l]
            end
        end
    end
    return C
end

function _mul(Bt::LinearAlgebra.Adjoint{Tv,Matrix{Tv}}, A::SparseCSC{Tv,Ti})::Matrix{Tv} where {Tv, Ti}
    B = Bt.parent
    k, m = size(B)
    k, n = size(A)
    C = zeros(m,n)
    for j = 1:n
        for z = A.colptr[j]:A.colptr[j+1]-1
            l = A.rowind[z]
            for i = 1:m
                C[i,j] += A.val[z] * B[l,i]
            end
        end
    end
    return C
end

function _mul(B::Matrix{Tv}, At::Adjoint{SparseCSC{Tv,Ti}})::Matrix{Tv} where {Tv, Ti}
    A = At.parent
    m, k = size(B)
    n, k = size(A)
    C = zeros(m,n)
    for l = 1:k
        for z = A.colptr[l]:A.colptr[l+1]-1
            j = A.rowind[z]
            for i = 1:m
                C[i,j] += A.val[z] * B[i,l]
            end
        end
    end
    return C
end

function _mul(Bt::LinearAlgebra.Adjoint{Tv,Matrix{Tv}}, At::Adjoint{SparseCSC{Tv,Ti}})::Matrix{Tv} where {Tv, Ti}
    A = At.parent
    B = Bt.parent
    k, m = size(B)
    n, k = size(A)
    C = zeros(m,n)
    for l = 1:k
        for z = A.colptr[l]:A.colptr[l+1]-1
            j = A.rowind[z]
            for i = 1:m
                C[i,j] += A.val[z] * B[l,i]
            end
        end
    end
    return C
end

#################### COO

function _mul(A::SparseCOO{Tv,Ti}, x::Vector{Tv})::Vector{Tv} where {Tv, Ti}
    m, n = size(A)
    y = zeros(m)
    for z = 1:SparseArrays.nnz(A)
        i = A.rowind[z]
        j = A.colind[z]
        y[i] += A.val[z] * x[j]
    end
    return y
end

function _mul(At::Adjoint{SparseCOO{Tv,Ti}}, x::Vector{Tv})::Vector{Tv} where {Tv, Ti}
    A = At.parent
    n, m = size(A)
    y = zeros(m)
    for z = 1:SparseArrays.nnz(A)
        j = A.rowind[z]
        i = A.colind[z]
        y[i] += A.val[z] * x[j]
    end
    return y
end

###

function _mul(A::SparseCOO{Tv,Ti}, B::Matrix{Tv})::Matrix{Tv} where {Tv, Ti}
    m, k = size(A)
    k, n = size(B)
    C = zeros(m,n)
    for z = 1:SparseArrays.nnz(A)
        i = A.rowind[z]
        l = A.colind[z]
        for j = 1:n
            C[i,j] += A.val[z] * B[l,j]
        end
    end
    return C
end

function _mul(A::SparseCOO{Tv,Ti}, Bt::LinearAlgebra.Adjoint{Tv,Matrix{Tv}})::Matrix{Tv} where {Tv, Ti}
    B = Bt.parent
    m, k = size(A)
    n, k = size(B)
    C = zeros(m,n)
    for z = 1:SparseArrays.nnz(A)
        i = A.rowind[z]
        l = A.colind[z]
        for j = 1:n
            C[i,j] += A.val[z] * B[j,l]
        end
    end
    return C
end

function _mul(At::Adjoint{SparseCOO{Tv,Ti}}, B::Matrix{Tv})::Matrix{Tv} where {Tv, Ti}
    A = At.parent
    k, m = size(A)
    k, n = size(B)
    C = zeros(m,n)
    for z = 1:SparseArrays.nnz(A)
        l = A.rowind[z]
        i = A.colind[z]
        for j = 1:n
            C[i,j] += A.val[z] * B[l,j]
        end
    end
    return C
end

function _mul(At::Adjoint{SparseCOO{Tv,Ti}}, Bt::LinearAlgebra.Adjoint{Tv,Matrix{Tv}})::Matrix{Tv} where {Tv, Ti}
    A = At.parent
    B = Bt.parent
    k, m = size(A)
    n, k = size(B)
    C = zeros(m,n)
    for z = 1:SparseArrays.nnz(A)
        l = A.rowind[z]
        i = A.colind[z]
        for j = 1:n
            C[i,j] += A.val[z] * B[j,l]
        end
    end
    return C
end

###

function _mul(B::Matrix{Tv}, A::SparseCOO{Tv,Ti})::Matrix{Tv} where {Tv, Ti}
    m, k = size(B)
    k, n = size(A)
    C = zeros(m,n)
    for z = 1:SparseArrays.nnz(A)
        l = A.rowind[z]
        j = A.colind[z]
        for i = 1:m
            C[i,j] += A.val[z] * B[i,l]
        end
    end
    return C
end

function _mul(Bt::LinearAlgebra.Adjoint{Tv,Matrix{Tv}}, A::SparseCOO{Tv,Ti})::Matrix{Tv} where {Tv, Ti}
    B = Bt.parent
    k, m = size(B)
    k, n = size(A)
    C = zeros(m,n)
    for z = 1:SparseArrays.nnz(A)
        l = A.rowind[z]
        j = A.colind[z]
        for i = 1:m
            C[i,j] += A.val[z] * B[l,i]
        end
    end
    return C
end

function _mul(B::Matrix{Tv}, At::Adjoint{SparseCOO{Tv,Ti}})::Matrix{Tv} where {Tv, Ti}
    A = At.parent
    m, k = size(B)
    n, k = size(A)
    C = zeros(m,n)
    for z = 1:SparseArrays.nnz(A)
        j = A.rowind[z]
        l = A.colind[z]
        for i = 1:m
            C[i,j] += A.val[z] * B[i,l]
        end
    end
    return C
end

function _mul(Bt::LinearAlgebra.Adjoint{Tv,Matrix{Tv}}, At::Adjoint{SparseCOO{Tv,Ti}})::Matrix{Tv} where {Tv, Ti}
    A = At.parent
    B = Bt.parent
    k, m = size(B)
    n, k = size(A)
    C = zeros(m,n)
    for z = 1:SparseArrays.nnz(A)
        j = A.rowind[z]
        l = A.colind[z]
        for i = 1:m
            C[i,j] += A.val[z] * B[l,i]
        end
    end
    return C
end

######## override

function Base.:*(x::Tv, A::AbstractSparseM{Tv,Ti}) where {Tv<:Number,Ti}
    B = copy(A)
    B.val .*= x
    return B
end

function Base.:*(A::AbstractSparseM{Tv,Ti}, x::Tv) where {Tv<:Number,Ti}
    B = copy(A)
    B.val .*= x
    return B
end

function Base.:/(A::AbstractSparseM{Tv,Ti}, x::Tv) where {Tv<:Number,Ti}
    B = copy(A)
    B.val ./= x
    return B
end

(Base.:*)(A::SparseCSR{Tv,Ti}, x::Vector{Tv}) where {Tv,Ti} = _mul(A, x)
(Base.:*)(A::Adjoint{SparseCSR{Tv,Ti}}, x::Vector{Tv}) where {Tv,Ti} = _mul(A, x)

(Base.:*)(A::SparseCSR{Tv,Ti}, B::Matrix{Tv}) where {Tv,Ti} = _mul(A, B)
(Base.:*)(A::SparseCSR{Tv,Ti}, B::LinearAlgebra.Adjoint{Tv,Matrix{Tv}}) where {Tv,Ti} = _mul(A, B)
(Base.:*)(A::Adjoint{SparseCSR{Tv,Ti}}, B::Matrix{Tv}) where {Tv,Ti} = _mul(A, B)
(Base.:*)(A::Adjoint{SparseCSR{Tv,Ti}}, B::LinearAlgebra.Adjoint{Tv,Matrix{Tv}}) where {Tv,Ti} = _mul(A, B)

(Base.:*)(B::Matrix{Tv}, A::SparseCSR{Tv,Ti}) where {Tv,Ti} = _mul(B, A)
(Base.:*)(B::LinearAlgebra.Adjoint{Tv,Matrix{Tv}}, A::SparseCSR{Tv,Ti}) where {Tv,Ti} = _mul(B, A)
(Base.:*)(B::Matrix{Tv}, A::Adjoint{SparseCSR{Tv,Ti}}) where {Tv,Ti} = _mul(B, A)
(Base.:*)(B::LinearAlgebra.Adjoint{Tv,Matrix{Tv}}, A::Adjoint{SparseCSR{Tv,Ti}}) where {Tv,Ti} = _mul(B, A)

###

(Base.:*)(A::SparseCSC{Tv,Ti}, x::Vector{Tv}) where {Tv,Ti} = _mul(A, x)
(Base.:*)(A::Adjoint{SparseCSC{Tv,Ti}}, x::Vector{Tv}) where {Tv,Ti} = _mul(A, x)

(Base.:*)(A::SparseCSC{Tv,Ti}, B::Matrix{Tv}) where {Tv,Ti} = _mul(A, B)
(Base.:*)(A::SparseCSC{Tv,Ti}, B::LinearAlgebra.Adjoint{Tv,Matrix{Tv}}) where {Tv,Ti} = _mul(A, B)
(Base.:*)(A::Adjoint{SparseCSC{Tv,Ti}}, B::Matrix{Tv}) where {Tv,Ti} = _mul(A, B)
(Base.:*)(A::Adjoint{SparseCSC{Tv,Ti}}, B::LinearAlgebra.Adjoint{Tv,Matrix{Tv}}) where {Tv,Ti} = _mul(A, B)

(Base.:*)(B::Matrix{Tv}, A::SparseCSC{Tv,Ti}) where {Tv,Ti} = _mul(B, A)
(Base.:*)(B::LinearAlgebra.Adjoint{Tv,Matrix{Tv}}, A::SparseCSC{Tv,Ti}) where {Tv,Ti} = _mul(B, A)
(Base.:*)(B::Matrix{Tv}, A::Adjoint{SparseCSC{Tv,Ti}}) where {Tv,Ti} = _mul(B, A)
(Base.:*)(B::LinearAlgebra.Adjoint{Tv,Matrix{Tv}}, A::Adjoint{SparseCSC{Tv,Ti}}) where {Tv,Ti} = _mul(B, A)

###

(Base.:*)(A::SparseCOO{Tv,Ti}, x::Vector{Tv}) where {Tv,Ti} = _mul(A, x)
(Base.:*)(A::Adjoint{SparseCOO{Tv,Ti}}, x::Vector{Tv}) where {Tv,Ti} = _mul(A, x)

(Base.:*)(A::SparseCOO{Tv,Ti}, B::Matrix{Tv}) where {Tv,Ti} = _mul(A, B)
(Base.:*)(A::SparseCOO{Tv,Ti}, B::LinearAlgebra.Adjoint{Tv,Matrix{Tv}}) where {Tv,Ti} = _mul(A, B)
(Base.:*)(A::Adjoint{SparseCOO{Tv,Ti}}, B::Matrix{Tv}) where {Tv,Ti} = _mul(A, B)
(Base.:*)(A::Adjoint{SparseCOO{Tv,Ti}}, B::LinearAlgebra.Adjoint{Tv,Matrix{Tv}}) where {Tv,Ti} = _mul(A, B)

(Base.:*)(B::Matrix{Tv}, A::SparseCOO{Tv,Ti}) where {Tv,Ti} = _mul(B, A)
(Base.:*)(B::LinearAlgebra.Adjoint{Tv,Matrix{Tv}}, A::SparseCOO{Tv,Ti}) where {Tv,Ti} = _mul(B, A)
(Base.:*)(B::Matrix{Tv}, A::Adjoint{SparseCOO{Tv,Ti}}) where {Tv,Ti} = _mul(B, A)
(Base.:*)(B::LinearAlgebra.Adjoint{Tv,Matrix{Tv}}, A::Adjoint{SparseCOO{Tv,Ti}}) where {Tv,Ti} = _mul(B, A)


function compute_orbital_energies(F, C)
  return transpose(C)*F*C
end


# Compute fast Walsh-Hadamard transforms (in natural, dyadic, or sequency order)
# using FFTW, by interpreting them as 2x2x2...x2x2 DFTs.  We follow Matlab's
# convention in which ifwht is the unnormalized transform and fwht has a 1/N
# normalization (as opposed to using a unitary normalization).
__precompile__(true)
module Hadamard
export fwht, ifwht, fwht_natural, ifwht_natural, fwht_dyadic, ifwht_dyadic, hadamard

using Base.FFTW
import Base.FFTW: set_timelimit, dims_howmany, unsafe_execute!, cFFTWPlan, r2rFFTWPlan, PlanPtr
import Base.DFT: normalization, complexfloat

using Compat

# A power-of-two dimension to be transformed is interpreted as a
# 2x2x2x....x2x2  multidimensional DFT.  This function transforms
# each individual power-of-two dimension n into the corresponding log2(n)
# DFT dimensions.  If bitreverse is true, then the output is in
# bit-reversed (transposed) order (which may not work in-place).
function hadamardize(dims::Array{Int,2}, bitreverse::Bool)
    ntot = 0
    for i = 1:size(dims,2)
        n = dims[1,i]
        if !ispow2(n)
            throw(ArgumentError("non power-of-two Hadamard-transform length"))
        end
        ntot += trailing_zeros(n)
    end
    hdims = Array(Int,3,ntot)
    j = 0
    for i = 1:size(dims,2)
        n = dims[1,i]
        is = dims[2,i]
        os = dims[3,i]
        log2n = trailing_zeros(n)
        krange = j+1:j+log2n
        for k in krange
            hdims[1,k] = 2
            hdims[2,k] = is
            hdims[3,k] = os
            is *= 2
            os *= 2
        end
        if bitreverse
            a = hdims[3,krange] # size = 1x3 in Julia 0.4, = 3 in 0.5
            hdims[3,krange] = flipdim(a, ndims(a))
        end
        j += log2n
    end
    return hdims
end

for (Tr,Tc,fftw,lib) in ((:Float64,:Complex128,"fftw",FFTW.libfftw),
                         (:Float32,:Complex64,"fftwf",FFTW.libfftwf))
    @eval function Plan_Hadamard{N}(X::StridedArray{$Tc,N}, Y::StridedArray{$Tc,N},
                                 region, flags::Unsigned, timelimit::Real,
                                 bitreverse::Bool)
        set_timelimit($Tr, timelimit)
        dims, howmany = dims_howmany(X, Y, [size(X)...], region)
        dims = hadamardize(dims, bitreverse)
        plan = ccall(($(string(fftw,"_plan_guru64_dft")),$lib),
                     PlanPtr,
                     (Int32, Ptr{Int}, Int32, Ptr{Int},
                      Ptr{$Tc}, Ptr{$Tc}, Int32, UInt32),
                     size(dims,2), dims, size(howmany,2), howmany,
                     X, Y, FFTW.FORWARD, flags)
        set_timelimit($Tr, NO_TIMELIMIT)
        if plan == C_NULL
            error("FFTW could not create plan") # shouldn't normally happen
        end
        return cFFTWPlan{$Tc,FFTW.FORWARD,X===Y,N}(plan, flags, region, X, Y)
    end

    @eval function Plan_Hadamard{N}(X::StridedArray{$Tr,N}, Y::StridedArray{$Tr,N},
                                 region, flags::Unsigned, timelimit::Real,
                                 bitreverse::Bool)
        set_timelimit($Tr, timelimit)
        dims, howmany = dims_howmany(X, Y, [size(X)...], region)
        dims = hadamardize(dims, bitreverse)
        kind = Array(Int32, size(dims,2))
        kind[:] = R2HC
        plan = ccall(($(string(fftw,"_plan_guru64_r2r")),$lib),
                     PlanPtr,
                     (Int32, Ptr{Int}, Int32, Ptr{Int},
                      Ptr{$Tr}, Ptr{$Tr}, Ptr{Int32}, UInt32),
                     size(dims,2), dims, size(howmany,2), howmany,
                     X, Y, kind, flags)
        set_timelimit($Tr, NO_TIMELIMIT)
        if plan == C_NULL
            error("FFTW could not create plan") # shouldn't normally happen
        end
        return r2rFFTWPlan{$Tr,(map(Int,kind)...),X===Y,N}(plan, flags, region, X, Y)
    end
end

############################################################################
# Define ifwht (inverse/unnormalized) transforms for various orderings

# Natural (Hadamard) ordering:

function ifwht_natural{T<:fftwNumber}(X::StridedArray{T}, region)
    Y = similar(X)
    p = Plan_Hadamard(X, Y, region, ESTIMATE, NO_TIMELIMIT, false)
    unsafe_execute!(p)
    return Y
end

function ifwht_natural{T<:Number}(X::StridedArray{T}, region)
    Y = float(X)
    p = Plan_Hadamard(Y, Y, region, ESTIMATE, NO_TIMELIMIT, false)
    unsafe_execute!(p)
    return Y
end

# Dyadic (Paley, bit-reversed) ordering:

function ifwht_dyadic{T<:fftwNumber}(X::StridedArray{T}, region)
    Y = similar(X)
    p = Plan_Hadamard(X, Y, region, ESTIMATE, NO_TIMELIMIT, true)
    unsafe_execute!(p)
    return Y
end

function ifwht_dyadic{T<:Number}(X::StridedArray{T}, region)
    return ifwht_dyadic(float(X), region)
end

############################################################################
# Sequency (Walsh) ordering:

# ifwht along a single dimension d of X
function ifwht{T<:fftwNumber}(X::Array{T}, region)
    Y = ifwht_dyadic(X, region)

    # Perform Gray-code permutation of Y (TODO: in-place?)
    if isempty(region)
        return Y
    elseif ndims(Y) == 1
        return [ Y[1 + ((i >> 1) $ i)] for i = 0:length(Y)-1 ]
    else
        sz = [size(Y)...]
        tmp = Array(T, maximum(sz[region])) # storage for out-of-place perm.
        for d in region
            # somewhat ugly loops to do 1d permutations along dimension d
            na = prod(sz[d+1:end])
            n = sz[d]
            nb = prod(sz[1:d-1])
            sa = nb * n
            for ia = 0:na-1
                for ib = 1:nb
                    i0 = ib + sa * ia
                    for i = 0:n-1
                        tmp[i+1] = Y[i0 + nb * ((i >> 1) $ i)]
                    end
                    for i = 0:n-1
                        Y[i0 + nb * i] = tmp[i+1]
                    end
                end
            end
        end
        return Y
    end
end

# handle 1d case of strided arrays (loops in multidim case are too annoying)
function ifwht{T<:fftwNumber}(X::StridedVector{T}, region)
    Y = ifwht_dyadic(X, region)

    # Perform Gray-code permutation of Y (TODO: in-place?)
    if isempty(region)
        return Y
    else
        return [ Y[1 + ((i >> 1) $ i)] for i = 0:length(Y)-1 ]
    end
end

# fallback for subarrays
function ifwht{T<:fftwNumber}(X::StridedArray{T}, region)
    return ifwht(copy(X), region)
end

# fallback for other types
function ifwht{T<:Number}(X::StridedArray{T}, region)
    return ifwht(float(X), region)
end

############################################################################
# Forward transforms (normalized by 1/N as in Matlab) and transforms
# without the region argument:

for f in (:ifwht_natural, :ifwht_dyadic, :ifwht)
    g = Symbol(string(f)[2:end])
    @eval begin
        $f(X) = $f(X, 1:ndims(X))
        $g(X) = scale!($f(X), normalization(X,1:ndims(X)))
        $g(X,r) = scale!($f(X,r), normalization(X,r))
    end
end

############################################################################
# Utilities to work with a precomputed cache of known Hadamard matrices
# of various sizes, produced by util/fetchhadamard.jl from Sloane's web page
# and stored as BitMatrices

function readcache(cachefile::AbstractString)
    B = BitMatrix[]
    open(cachefile, "r") do io
        while !eof(io)
            k = convert(Int, ntoh(read(io, Int64)))
            b = BitArray(k, k)
            # Hack: use internal binary data from BitArray for efficiency
            bits = read(io, eltype(b.chunks), length(b.chunks))
            for i = 1:length(bits)
                b.chunks[i] = ntoh(bits[i])
            end
            push!(B, b)
        end
    end
    return B
end
readcache() = readcache(joinpath(dirname(@__FILE__), "..", "src", "cache.dat"))

function printsigns{T<:Real}(io::IO, A::AbstractMatrix{T})
    m, n = size(A)
    println(io, m, "x", n, " sign matrix from ", typeof(A))
    for i = 1:m
        for j = 1:n
            print(io, A[i,j] > 0 ? "+" : "-")
        end
        println(io)
    end
end
printsigns(A) = printsigns(STDOUT, A)

function frombits(B::BitMatrix)
    A = convert(Matrix{Int8}, B)
    for i = 1:length(A)
        A[i] = A[i] == 0 ? -1 : 1
    end
    return A
end

const hadamards = BitMatrix[] # read lazily below

############################################################################
# Create Int8 order-n Hadamard matrix (in natural order), by factorizing
# n into a product of known Hadamard sizes (if possible)

const H2 = Int8[1 1; 1 -1]

function hadamard(n::Integer)
    n < 1 && throw(ArgumentError("size n=$n should be positive"))
    n0 = n
    H = reshape(Int8[1], 1,1)
    if !ispow2(n)
        isempty(hadamards) && append!(hadamards, readcache())
        for i = length(hadamards):-1:1
            k = size(hadamards[i], 1)
            if rem(n, k) == 0
                Hk = frombits(hadamards[i])
                while true
                    H = kron(H, Hk)
                    n = div(n, k)
                    rem(n, k) != 0 && break
                end
            end
        end
    end
    if !ispow2(n)
        n >>= trailing_zeros(n)
        throw(ArgumentError("unknown Hadamard factor $n of $n0"))
    end
    # Note: it would be faster to do a "power-by-squaring" like algorithm
    # here, where we repeatedly double the order via kron(H, H), but
    # it's not clear to me that we care much about performance here.
    while n > 1
        H = kron(H, H2)
        n >>= 1
    end
    return H
end

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

end # module


# MultiVector arithmetic for CliffordAlgebras.jl

import Base.+, Base.-, Base.*, Base./, Base.\
import Base.inv, Base.adjoint, Base.exp
import LinearAlgebra.norm, LinearAlgebra.norm_sqr
import LinearAlgebra.SingularException
import StaticArrays.SVector, StaticArrays.SMatrix
import SparseArrays.sparse, SparseArrays.sparsevec

@generated function (+)(a::MultiVector{CA,Ta}, b::MultiVector{CA,Tb}) where {CA,Ta,Tb}
    d = dimension(CA)
    acc_a = [Int(0) for n = 1:d]
    acc_b = [Int(0) for n = 1:d]
    bi_a = baseindices(a)
    bi_b = baseindices(b)
    acc_a[collect(bi_a)] = 1:length(bi_a)
    acc_b[collect(bi_b)] = 1:length(bi_b)
    BI = Tuple(findall(map( (a,b) -> !iszero(a) || !iszero(b), acc_a, acc_b)))
    K = length(BI)
    T = promote_type(Ta, Tb)
    if iszero(K)
        return :(MultiVector(Algebra(a), zero($T)))
    end
    coeffs = (
        iszero(acc_a[n]) ? :(coefficients(b)[$(acc_b[n])]) :
        iszero(acc_b[n]) ? :(coefficients(a)[$(acc_a[n])]) : :(coefficients(a)[$(acc_a[n])] + coefficients(b)[$(acc_b[n])]) for
        n in BI
    )
    coeffsexpr = Expr(:call, :(NTuple{$K,$T}), Expr(:call, :tuple, coeffs...))
    :(@inbounds MultiVector(Algebra(a), $BI, $coeffsexpr))
end

(+)(a::Real, b::MultiVector{CA}) where {CA} = MultiVector(CA, a) + b
(+)(a::MultiVector{CA}, b::Real) where {CA} = a + MultiVector(CA, b)

(-)(a::MultiVector) = -1 * a
(-)(a::MultiVector{CA}, b::MultiVector{CA}) where {CA} = a + (-b)
(-)(a::Real, b::MultiVector) = a + (-b)
(-)(a::MultiVector, b::Real) = a + (-b)

function generatefilteredprod(
    a::Type{<:MultiVector{CA,Ta}},
    b::Type{<:MultiVector{CA,Tb}},
    ::Val{Filter},
) where {CA,Ta,Tb,Filter}
    d = dimension(CA)
    acc = [NTuple{3,Int}[] for n = 1:d]
    for (ka, base_a) in enumerate(baseindices(a)), (kb, base_b) in enumerate(baseindices(b))
        baseidx, coeff = baseproduct(CA, base_a, base_b)
        leftgrade = basegrade(CA, base_a)
        rightgrade = basegrade(CA, base_b)
        productgrade = basegrade(CA, baseidx)
        keep = if Filter == :geometric
            true
        elseif Filter == :exterior
            productgrade == leftgrade + rightgrade
        elseif Filter == :rightcontraction
            productgrade == leftgrade - rightgrade
        elseif Filter == :leftcontraction
            productgrade == rightgrade - leftgrade
        elseif Filter == :fatdot
            productgrade == abs(leftgrade - rightgrade)
        elseif Filter == :scalar
            iszero(productgrade)
        else
            error("unknown filter.")
        end
        if !iszero(coeff) && keep
            push!(acc[baseidx], (ka, kb, coeff))
        end
    end
    BI = Tuple(findall(!isempty, acc))
    K = length(BI)
    T = promote_type(Ta, Tb)
    if iszero(K)
        return :(MultiVector(Algebra(a), zero($T)))
    end
    coeffs = (
        Expr(
            :call,
            :+,
            (:($(tup[3]) * coefficients(a)[$(tup[1])] * coefficients(b)[$(tup[2])]) for tup in acc[n])...,
        ) for n in BI
    )
    coeffsexpr = Expr(:call, :(NTuple{$K,$T}), Expr(:call, :tuple, coeffs...))
    :(@inbounds MultiVector(Algebra(a), $BI, $coeffsexpr))
end

@generated function geometricprod(
    a::MultiVector{CA,Ta},
    b::MultiVector{CA,Tb},
) where {CA,Ta,Tb}
    generatefilteredprod(a,b,Val(:geometric))
end

@generated function exteriorprod(
    a::MultiVector{CA,Ta},
    b::MultiVector{CA,Tb},
) where {CA,Ta,Tb}
    generatefilteredprod(a,b,Val(:exterior))
end

@generated function leftcontractionprod(
    a::MultiVector{CA,Ta},
    b::MultiVector{CA,Tb},
) where {CA,Ta,Tb}
    generatefilteredprod(a,b,Val(:leftcontraction))
end

@generated function rightcontractionprod(
    a::MultiVector{CA,Ta},
    b::MultiVector{CA,Tb},
) where {CA,Ta,Tb}
    generatefilteredprod(a,b,Val(:rightcontraction))
end

@generated function fatdotprod(
    a::MultiVector{CA,Ta},
    b::MultiVector{CA,Tb},
) where {CA,Ta,Tb}
    generatefilteredprod(a,b,Val(:fatdot))
end

@generated function scalarprod(
    a::MultiVector{CA,Ta},
    b::MultiVector{CA,Tb},
) where {CA,Ta,Tb}
    generatefilteredprod(a,b,Val(:scalar))
end

"""
    a * b

Calculates the geometric product of two MultiVectors a and b.
"""
(*)(a::MultiVector{CA}, b::MultiVector{CA}) where {CA} = geometricprod(a, b)
(*)(a::Real, b::MultiVector) = MultiVector(Algebra(b), baseindices(b), a .* coefficients(b))
(*)(a::MultiVector, b::Real) = b * a

"""
    a ∧ b

Calculates the wedge product between two MultiVectors a and b.
"""
(∧)(a::MultiVector{CA}, b::MultiVector{CA}) where {CA} = exteriorprod(a, b)
(∧)(a::MultiVector{CA}, b::Real) where {CA} = a ∧ MultiVector(CA,b)
(∧)(a::Real, b::MultiVector{CA}) where {CA} = MultiVector(CA,a) ∧ b
(∧)(a::Real, b::Real) = a * b

"""
    a ⋅ b

Calculates the "fat dot" product between the MultiVectors a and b.
"""
(⋅)(a::MultiVector{CA}, b::MultiVector{CA}) where {CA} = fatdotprod(a, b)
(⋅)(a::MultiVector{CA}, b::Real) where {CA} = a ⋅ MultiVector(CA,b)
(⋅)(a::Real, b::MultiVector{CA}) where {CA} = MultiVector(CA,a) ⋅ b
(⋅)(a::Real, b::Real) = a * b

"""
    a ⨼ b

Calculates the left contraction of the MultiVectors a and b.
"""
(⨼)(a::MultiVector{CA}, b::MultiVector{CA}) where {CA} = leftcontractionprod(a, b)
(⨼)(a::MultiVector{CA}, b::Real) where {CA} = a ⨼ MultiVector(CA,b)
(⨼)(a::Real, b::MultiVector{CA}) where {CA} = MultiVector(CA,a) ⨼ b
(⨼)(a::Real, b::Real) = a * b

"""
    a ⨽ b

Calculates the right contraction of the MultiVectors a and b.
"""
(⨽)(a::MultiVector{CA}, b::MultiVector{CA}) where {CA} = rightcontractionprod(a, b)
(⨽)(a::MultiVector{CA}, b::Real) where {CA} = a ⨽ MultiVector(CA,b)
(⨽)(a::Real, b::MultiVector{CA}) where {CA} = MultiVector(CA,a) ⨽ b
(⨽)(a::Real, b::Real) = a * b

"""
    a ⋆ b

Calculates the scalar product of the MultiVectors a and b.
"""
(⋆)(a::MultiVector{CA}, b::MultiVector{CA}) where {CA} = scalarprod(a, b)
(⋆)(a::MultiVector{CA}, b::Real) where {CA} = a ⋆ MultiVector(CA,b)
(⋆)(a::Real, b::MultiVector{CA}) where {CA} = MultiVector(CA,a) ⋆ b
(⋆)(a::Real, b::Real) = a * b

"""
    a ∨ b

Calculates the regressive product of the MultiVectors a and b.
"""
(∨)(a::MultiVector{CA}, b::MultiVector{CA}) where {CA} = dual(dual(a) ∧ dual(b))
(∨)(a::MultiVector{CA}, b::Real) where {CA} = a ∨ MultiVector(CA,b)
(∨)(a::Real, b::MultiVector{CA}) where {CA} = MultiVector(CA,a) ∨ b
(∨)(a::Real, b::Real) where {CA} = zero(promote_type(typeof(a),typeof(b)))

function generatesymprod(
    a::Type{<:MultiVector{CA,Ta}},
    b::Type{<:MultiVector{CA,Tb}},
    ::Val{P}
) where {CA,Ta,Tb,P}
    d = dimension(CA)
    acc = [NTuple{3,Int}[] for n = 1:d]
    for (ka, base_a) in enumerate(baseindices(a)), (kb, base_b) in enumerate(baseindices(b))
        baseidx1, coeff1 = baseproduct(CA, base_a, base_b)
        baseidx2, coeff2 = baseproduct(CA, base_b, base_a)
        if baseidx1 == baseidx2
            coeff = coeff1 + P * coeff2
            if !iszero(coeff)
                push!(acc[baseidx1], (ka, kb, coeff))
            end
        else
            if !iszero(coeff1)
                push!(acc[baseidx1], (ka, kb, coeff1))
            end
            if !iszero(coeff2)
                push!(acc[baseidx2], (ka, kb, P * coeff2))
            end
        end
    end
    BI = Tuple(findall(!isempty, acc))
    K = length(BI)
    T = promote_type(Ta, Tb)
    if iszero(K)
        return :(MultiVector(Algebra(a), zero($T)))
    end
    coeffs = (
        Expr(
            :call,
            :+,
            (:($(tup[3]) * coefficients(a)[$(tup[1])] * coefficients(b)[$(tup[2])]) for tup in acc[n])...,
        ) for n in BI
    )
    coeffsexpr = Expr(:call, :(NTuple{$K,$T}), Expr(:call, :tuple, coeffs...))
    :(@inbounds MultiVector(Algebra(a), $BI, $coeffsexpr))
end

@generated function commutatorprod(
    a::MultiVector{CA,Ta},
    b::MultiVector{CA,Tb}
) where {CA,Ta,Tb}
    generatesymprod(a, b, Val(-1))
end

@generated function anticommutatorprod(
    a::MultiVector{CA,Ta},
    b::MultiVector{CA,Tb}
) where {CA,Ta,Tb}
    generatesymprod(a, b, Val(+1))
end

"""
    a ×₋ b

Calculates the commutator ab-ba of two MultiVectors a and b.
"""
(×₋)(a::MultiVector{CA,Ta}, b::MultiVector{CA,Tb}) where {CA,Ta,Tb} =
    inv(promote_type(Ta, Tb)(2)) * commutatorprod(a, b)

(×₋)(::Ta, ::MultiVector{CA,Tb}) where {CA,Ta<:Real,Tb} = zero(promote_type(Ta, Tb))
(×₋)(a::MultiVector, b::Real) = b ×₋ a
(×₋)(::Ta, ::Tb) where {Ta<:Real,Tb<:Real} = zero(promote_type(Ta, Tb))

"""
    a ×₊ b

Calculates the anti-commutator ab+ba of two MultiVectors a and b.
"""
(×₊)(a::MultiVector{CA,Ta}, b::MultiVector{CA,Tb}) where {CA,Ta,Tb} =
    inv(promote_type(Ta, Tb)(2)) * anticommutatorprod(a, b)
(×₊)(a::Real, b::MultiVector) = a * b
(×₊)(a::MultiVector, b::Real) = a * b
(×₊)(a::Real, b::Real) = a * b

@generated function sandwichproduct(
    a::MultiVector{CA,Ta},
    b::MultiVector{CA,Tb},
) where {CA,Ta,Tb}
    # a * b * reverse(a)
    d = dimension(CA)
    acc = [NTuple{4,Int}[] for n = 1:d]
    for (ka, base_a) in enumerate(baseindices(a)),
        (kb, base_b) in enumerate(baseindices(b)),
        (ka_rev, base_a_rev) in enumerate(baseindices(a))

        leftbaseidx, leftcoeff = baseproduct(CA, base_a, base_b)
        baseidx, rightcoeff = baseproduct(CA, leftbaseidx, base_a_rev)
        coeff = leftcoeff * rightcoeff * basereverse(CA, base_a_rev)
        if !iszero(coeff)
            push!(acc[baseidx], (ka, kb, ka_rev, coeff))
        end
    end
    BI = Tuple(findall(!isempty, acc))
    K = length(BI)
    T = promote_type(Ta, Tb)
    if iszero(K)
        return :(MultiVector(Algebra(a), zero($T)))
    end
    coeffs = (
        Expr(
            :call,
            :+,
            (
                :($(tup[4]) * coefficients(a)[$(tup[1])] * coefficients(b)[$(tup[2])] * coefficients(a)[$(tup[3])]) for
                tup in acc[n]
            )...,
        ) for n in BI
    )
    coeffsexpr = Expr(:call, :(NTuple{$K,$T}), Expr(:call, :tuple, coeffs...))
    :(@inbounds MultiVector(Algebra(a), $BI, $coeffsexpr))
end

"""
    a ≀ b

Calculates the sandwich product a*b*reverse(a) for two MultiVectors a and b.
"""
(≀)(a::MultiVector{CA}, b::MultiVector{CA}) where {CA} = sandwichproduct(a, b)
(≀)(a::Real, b::MultiVector{CA}) where {CA} = a^2 * b
(≀)(a::MultiVector{CA}, b::Real) where {CA} = b * a *reverse(a)
(≀)(a::Real, b::Real) = a^2 * b

"""
    polarize(mv::MultiVector)
    mv'

Calculates the polarization of the MultiVector, i.e. mv * pseudoscalar.
"""
polarize(mv::MultiVector{CA}) where {CA} = mv * pseudoscalar(CA)

adjoint(mv::MultiVector) = polarize(mv)

"""
    Λᵏ(::MultiVector, ::Val{k}) where k
    Λᵏ(::MultiVector, k::Integer)

Projects the MultiVector onto k-vectors. Similar to grade(mv,k), but uses
@generated code and compile time optimizations.
"""
@generated function (Λᵏ)(mv::MultiVector, ::Val{k}) where {k}
    CA = algebra(mv)
    indexbounds =
        1 .+ (0, cumsum(ntuple(i -> binomial(order(CA), i - 1), 1 + order(CA)))...)
    @assert 0 < k + 1 < length(indexbounds) "k out of bounds."
    bi_low = indexbounds[k+1]
    bi_high = indexbounds[k+2] - 1
    s = findall(i -> i >= bi_low && i <= bi_high, baseindices(mv))
    BI = baseindices(mv)[s]
    K = length(BI)
    T = eltype(mv)
    coeffs = (:(coefficients(mv)[$i]) for i in s)
    if length(coeffs) == 0
        return :(MultiVector(Algebra(a), zero($T)))
    end
    coeffsexpr = Expr(:call, :(NTuple{$K,$T}), Expr(:call, :tuple, coeffs...))
    :(@inbounds MultiVector(Algebra(mv), $BI, $coeffsexpr))
end

(Λᵏ)(mv::MultiVector, k::Integer) = Λᵏ(mv, Val(k))

"""
    norm(::MultiVector)

Calculates the MultiVector norm defined as sqrt(grade(mv*reverse(mv),0))
"""
norm(mv::MultiVector) = sqrt(norm_sqr(mv))

"""
    norm_sqr(::MultiVector)

Calculates the MultiVector squared norm defined as grade(mv*reverse(mv),0)
"""
norm_sqr(mv::MultiVector) = scalar(Λᵏ(mv * reverse(mv), 0))

"""
    inv(::MultiVector)

Finds the inverse of the MultiVector. If no inverse exists a SingularException is thrown.
"""
@generated function inv(mv::MultiVector)
#function inv_generator(mv::Type{<:MultiVector})
    # Strategy: We can reduce the size of the problem
    # by looking only at self-adjoint multivectors.
    # Note that inv(a) == a' * inv(a' a) == a' * inv(a a').
    # Both a a' and a' a are self-adjoint with respect to reversal and
    # are therefore in a lower dimensional subspace of the algebra.
    # If b is self-adjoint, then so is inv(b) because
    # b * inv(b) = 1 = (1)' = (inv(b) * b)' = b' * inv(b)' = b * inv(b)'
    # => inv(b) == inv(b)'
    # A multivector that is self-adjoint also has a basis expansion that
    # consists only of self-adjoint basis vectors.

    CA = algebra(mv)
    T = eltype(mv)
    BI = baseindices(mv)
    d = dimension(CA)
    basesquares = map(n -> baseproduct(CA, n, n)[2], BI)
    reversesigns = map(n -> basereverse(CA, n), BI)
    if all(iszero, basesquares)
        # singular in all cases
        return :(throw(SingularException(0)))
    end
    if length(BI) == 1
        # trivial inverse
        rc = inv(basesquares[1])
        return quote
            c = coefficients(mv)[1]
            if iszero(c)
                throw(SingularException(1))
            end
            MultiVector(Algebra(mv), $BI, ($rc * inv(c),))
        end
    end
    if all(isone, reversesigns) || all(s -> isone(-s), reversesigns)
        # We're in the (anti-)self-adjoint subspace already
        BIinv = Tuple(findall(isequal(reversesigns[1]), ntuple(i -> basereverse(CA, i), d)))
        BIone = Tuple(
            sort(
                unique(
                    skipmissing(
                        Iterators.flatten([
                            iszero(baseproduct(CA, r, c)[2]) ? missing : baseproduct(CA, r, c)[1] for r in BI,
                            c in BIinv
                        ]),
                    ),
                ),
            ),
        )
        ea = [Expr(:call, :+) for r = 1:length(BIone), c = 1:length(BIinv)]
        for (i1, bi1) in enumerate(BI), (i2, bi2) in enumerate(BIinv)
            (bi3, pc) = baseproduct(CA, bi1, bi2)
            i3 = findfirst(isequal(bi3), BIone)
            push!(ea[i3, i2].args, :($pc * c[$i1]))
        end
        if d <= 8   # this is just a guess and needs some measurements
            ea = map(ex -> length(ex.args) == 1 ? :(zero($T)) : ex, ea)
            ma = Expr(:call, Expr(:curly, :SMatrix, length(BIone), length(BIinv)), ea[:]...)
            id = Expr(:call, :SVector, ntuple(i -> i == 1 ? one(T) : zero(T), length(BIone))...)
            return quote
                c = coefficients(mv)
                ma = $ma
                id = $id
                cinv = ma \ id
                if all(isapprox.(ma * cinv - id, zero($T); atol = 1e-8))
                    prune(MultiVector(Algebra(mv), $BIinv, Tuple(cinv)))
                else
                    throw(SingularException(2))
                end
            end
        else
            iT = typeof(inv(one(T)))
            nzea = findall(ex -> length(ex.args) > 1, ea)
            rows = map( ci -> ci[1], nzea)
            cols = map( ci -> ci[2], nzea)
            vals = map( ci -> ea[ci], nzea)
            rowsexpr = Expr(:ref, :Int, rows...)
            colsexpr = Expr(:ref, :Int, cols...)
            valsexpr = Expr(:ref, :($iT), vals... )
            rowcount = length(BIone)
            colcount = length(BIinv)
            return quote
                c = coefficients(mv)
                rows = $rowsexpr
                cols = $colsexpr
                vals = $valsexpr
                ma = sparse(rows, cols, vals, $rowcount, $colcount)
                id = sparsevec([1],[one($iT)],$rowcount)
                cinv = Array(ma) \ Array(id)
                if all(isapprox.(ma * cinv - id, zero($iT); atol = 1e-8))
                    prune(MultiVector(Algebra(mv), $BIinv, Tuple(cinv)))
                else
                    throw(SingularException(3))
                end
            end
        end
    end
    # We reduce the general case to that of self-adjoint inverses
    products = Tuple(
        (i1, i2, baseproduct(CA, b1, b2)...) for
        (i1, b1) in enumerate(BI), (i2, b2) in enumerate(BI) if
        !iszero(baseproduct(CA, b1, b2)[2]) && isone(basereverse(CA, baseproduct(CA, b1, b2)[1]))
    )
    BIsa = Tuple(sort(unique(map(p -> p[3], products))))
    coeffexpr = [Expr(:call, :+) for b in BIsa]
    for (i1, i2, bres, c) in products
        rc = basereverse(CA, BI[i2])
        push!(
            coeffexpr[findfirst(isequal(bres), BIsa)].args,
            :($(c * rc) * coefficients(mv)[$i1] * coefficients(mv)[$i2]),
        )
    end
    K = length(BIsa)
    coeffexpr = Expr(:call, Expr(:curly, :NTuple, K, T), Expr(:call, :tuple, coeffexpr...))
    quote
        @inbounds reverse(mv) * inv(MultiVector(Algebra(mv), $BIsa, $coeffexpr))
    end
end

function inv_fallback(mv::MultiVector)
    M = matrix(mv)
    e = vector(basevector(algebra(mv), :𝟏))
    c = M \ e
    bi = findall(!iszero, c)
    MultiVector(Algebra(mv), Tuple(bi), Tuple(c[bi]))
end

"""
    a / b
    (/)(a::MultiVector{CA}, b::MuliVector{CA}) where CA

Calculates the MultiVector quotient a/b by evaluating a*inv(b).
"""
(/)(a::MultiVector{CA}, b::MultiVector{CA}) where {CA} = a * inv(b)
(/)(a::Real, b::MultiVector{CA}) where {CA} = a * inv(b)
(/)(a::MultiVector{CA}, b::Real) where {CA} = a * inv(b)

"""
    b \\ a
    (\\)(b::MultiVector{CA}, b::MuliVector{CA}) where CA

Calculates the MultiVector quotient a/b by evaluating inv(b)*a.
"""
(\)(a::MultiVector{CA}, b::MultiVector{CA}) where {CA} = inv(a) * b
(\)(a::Real, b::MultiVector{CA}) where {CA} = inv(a) * b
(\)(a::MultiVector{CA}, b::Real) where {CA} = inv(a) * b

"""
    exp(::MultiVector)

Calculates the exponential function of the MultiVector defined by analytic continuation.
The generated code is automaticall specialized for the sparse representation of the MultiVector.
It may take advantage of commuting base vectors and split off exponential factors. Hyperbolic, trigonometric and nilpotent solutions are recognized and handled separately.
Calling prune or grade before exp may help to find the best algorithm for the exponential evaluation.
"""
@generated function exp(mv::MultiVector)
    # 1. Identify (optimal) subsets of base vectors so that all vectors in such a set commute with all vectors outside the set.
    ca = algebra(mv)
    bi = baseindices(mv)
    ct = [iszero(baseproduct(ca, i, j)[2] - baseproduct(ca, j, i)[2]) for i in bi, j in bi]
    ncsets = Vector{Int}[]
    for (ki, i) in enumerate(bi)
        nc = union(collect(bi[findall(!, ct[ki, :])]), Int[i])
        si = findfirst(g -> !isempty(intersect(g, nc)), ncsets)
        if isnothing(si)
            push!(ncsets, nc)
        else
            ncsets[si] = union(nc, ncsets[si])
        end
    end
    ncsets = map(sort, ncsets)
    prodexpr = Expr(:call, :*, :(one(eltype(mv))))
    for ncset in ncsets
        if length(ncset) == 1 && ncset[1] == 1
            push!(prodexpr.args, :(exp(scalar(mv))))
        else
            baseselector = Tuple([findfirst(isequal(i), bi) for i in ncset])
            coeftuple = Expr(:call, :tuple, map(i -> :(coefficients(mv)[$i]), baseselector)...)
            ncsetmvexpr = :(MultiVector(Algebra(mv), $(Tuple(ncset)), $coeftuple))
            basesquares = [baseproduct(ca, i, i)[2] for i in ncset]
            crosstermscancel =
                all(iszero(baseproduct(ca, i, j)[2] + baseproduct(ca, j, i)[2]) for i in ncset, j in ncset if i < j)
            if all(basesquares .== basesquares[1]) && crosstermscancel
                # We can simplify the series expansion
                if isone(basesquares[1])
                    # hyperbolic
                    push!(prodexpr.args, :(exp_hyp($ncsetmvexpr)))
                elseif isone(-basesquares[1])
                    # trigonometric
                    push!(prodexpr.args, :(exp_trig($ncsetmvexpr)))
                elseif iszero(basesquares[1])
                    # affine
                    push!(prodexpr.args, :(one(eltype(mv)) + $ncsetmvexpr))
                else
                    error("Internal error. Expected base vector to square to -1, 0 or 1")
                end
            else
                # brute force series calculation required
                push!(prodexpr.args, :(exp_fallback2($ncsetmvexpr)))
            end
        end
    end
    :(@inbounds $prodexpr)
end

function exp_trig(mv::MultiVector)
    if iszero(mv)
        return one(mv)
    else
        x = sqrt(sum(coefficients(mv) .^ 2))
        return cos(x) + sin(x) * mv * inv(x)
    end
end

function exp_hyp(mv::MultiVector)
    if iszero(mv)
        return one(mv)
    else
        x = sqrt(sum(coefficients(mv) .^ 2))
        return cosh(x) + sinh(x) * mv * inv(x)
    end
end

function exp_fallback(mv::MultiVector; order::Integer = 32)
    result = one(mv)
    xpower = one(mv)
    c = inv(one(eltype(mv)))
    for n = 1:order
        c /= n
        xpower *= mv
        result += c * xpower
    end
    result
end

function exp_fallback2(mv::MultiVector; order::Integer = 16)
    s = inv(2.0^order)
    x = mv * s
    h = 1 + x * (x * (x * inv(6) + inv(2)) + 1)
    for n = 1:order
        h *= h
    end
    h
end

"""
    outermorphism(A::AbstractMatrix, mv::MultiVector)

Calculates the outermorphism f of the MultiVector defined by f(v) = Av if v is in the grade-1 subspace of the algebra.
"""
@generated function outermorphism(A::AbstractMatrix, mv::MultiVector)
    ca = Algebra(mv)
    bi = baseindices(mv)
    bt = basetable(ca)
    bv = reduce(union, bt[collect(bi)])
    if isempty(bv)
        return :(mv)
    else
        nb = Expr(:call,:tuple)
        append!(nb.args, [:(MultiVector(ca, Tuple(A[:,$(bv[i])]))) for i = 1:length(bv)])
        s = Expr(:call,:+)
        for (k,b) in enumerate(bt[collect(bi)])
            ft = map(n -> findfirst(isequal(n), bv) ,b)
            if isempty(ft)
                push!(s.args, :(coefficients(mv)[$k]))
            elseif length(ft) == 1
                push!(s.args, :(coefficients(mv)[$k] * f[$(ft[1])]))
            else
                bp = foldl((a,b) -> Expr(:call, :∧, a, b), map(n->:(f[$n]), ft))
                push!(s.args, :(coefficients(mv)[$k] * $bp))
            end
        end
        quote
            @assert size(A,1) == size(A,2) == $(order(ca)) "The matrix A must be N×N where N is the order of the algebra."
            ca = Algebra(mv)
            f = $nb
            @inbounds $s
        end
    end
end


using SimpleDrawing, Plots, Clines
#resize_gr_window()

C = Circle(0, 0, 15)
D = Circle(23, 0, 8)

newdraw()
draw(C, linewidth = 2);
draw(D, linewidth = 2);

x = 15 * cos(pi / 6)
y = 15 * sin(pi / 6)

draw_vector(x, y, linecolor = :black)
annotate!(x / 2, y / 2 + 1, "a")

x = 8 * cos(pi / 6)
y = 8 * sin(pi / 6)
draw_vector(x, y, 23, 0, color = :black)

annotate!(x + 19, y / 2 + 0.7, "b")

#draw_point(23+offset,0,marker=4,color=:red,linecolor=:white)

offset = 6
draw(Circle(23 + offset + 0.2, 0, 0.2), fill = true, linecolor = :red, color = :red)

draw_vector(offset, 0, 23, 0, linecolor = :red)
annotate!(26, -0.5, "offset", annotationfontsize = 10)

draw_point(0, 0, color = :black, marker = 4)
draw_point(23, 0, color = :black, marker = 4)

finish()


#__precompile__()

module ExperimentalDesign

using Random
using Logging
using Primes
using DataFrames
using Distributions
using StatsModels
using DocStringExtensions
using LinearAlgebra

import Random: rand

# Types

export    AbstractDesign,   AbstractScreeningDesign,    AbstractFactorialDesign,
    AbstractOptimalDesign,  PlackettBurman, FullFactorial,  FractionalFactorial,
    OptimalDesign, DesignDistribution, CategoricalFactor

# Methods

export fullfactorial,  explicit_fullfactorial, plackettburman, fold!, isplackettburman,
    paley, rand, random_design!, next_offset_divisible_prime, kl_exchange, d_criterion

# Pretty printing

export show

include("design.jl")
include("plackett_burman.jl")
include("factorial.jl")
include("random.jl")
include("kl_exchange.jl")
include("custom_show.jl")
include("categorical.jl")

end


include("Tree.jl");

x3 = Variable(UInt(3), semantic=[1.,4.,-2.]);

tree = DivNode(
    Constant(2.) ,
    AddNode(
        DivNode(
            Constant(3.) ,
            SubNode(
                Constant(1.) ,
                DivNode(
                    Constant(2.),
                    x3
                )
            )
        ) ,
        Constant(1.)
    )
);

println("Figure 1: ", string(tree));

targets = [5.,4.,1.];
evaluateTree(tree);

initialEquation = NodeEquation(1., targets, 0., -1.);

calculateEquations!(tree,initialEquation);

println("Table 1:");

nodes, = iterateTree(tree);
for numNode in 1:length(nodes)
    println("-----------------------------------------------------------------------------------------")
    if (isa(nodes[numNode], NonTerminal))
        println("Node ", numNode, ": ", nodes[numNode].name);
    else
        println("Node ", numNode, ": ", string(nodes[numNode]));
    end;
    showEquation(nodes[numNode].equation, nodes[numNode].semantic);
end;

println("-----------------------------------------------------------------------------------------")
println("-----------------------------------------------------------------------------------------")
println("-----------------------------------------------------------------------------------------")

x1 = Variable(UInt(1), semantic=[3.,2.,0.]);
x2 = Variable(UInt(2), semantic=[0.,1.,0.]);
x3 = Variable(UInt(3), semantic=[1.,0.,3.]);
# x4 = Variable(UInt(4), semantic=[2.,0.,4.]);

tree = DivNode(
    Constant(2.) ,
    AddNode(
        DivNode(
            x1,
            Constant(2.)
        ) ,
        MulNode(
            x2,
            SubNode(
                MulNode(
                    x3,
                    # DivNode(
                        Constant(2.) ,
                        # x4
                    # ) ,
                ) ,
                Constant(1.)
            )
        )
    )
);
println("Figure 3: ", string(tree));

evaluateTree(tree);

targets = [5.,4.,1.];
initialEquation = NodeEquation(1., targets, 0., -1.);

calculateEquations!(tree,initialEquation);

println("Table 3: ");

nodes, = iterateTree(tree);
for numNode in 1:length(nodes)
    println("-----------------------------------------------------------------------------------------")
    if (isa(nodes[numNode], NonTerminal))
        println("Node ", numNode, ": ", nodes[numNode].name);
    else
        println("Node ", numNode, ": ", string(nodes[numNode]));
    end;
    showEquation(nodes[numNode].equation, nodes[numNode].semantic);
end;


# This file performs computations for rendering the Julia set iteration
#    z -> z - (z^3-1)/(3z^2)
# It takes a number of shortcuts and makes some optimizations for faster
# performance---this should not be viewed as definitive for computation
# of fractals.

const ITER = 15
const ITER_wait = 5
const sthresh = 5.0

function iterate!{T<:Real}(s::Matrix{T}, z_r::Matrix{T}, z_i::Matrix{T}, xlim, ylim)
    fill!(s, 0)
    one_third = convert(T, 1/3)
    @inbounds for (iy,yv) in enumerate(linspace(ylim[1],ylim[2],size(s,2)))
        for (ix,xv) in enumerate(linspace(xlim[1],xlim[2],size(s,1)))
            z_r[ix,iy] = xv
            z_i[ix,iy] = yv
        end
        for n = 1:ITER
            sf = n > ITER_wait ? 1 : 0
            @simd for ix = 1:size(s,1)
                # z -> z - (z^3-1)/(3*z^2) iteration, optimized for performance
                # The complex inversion is implemented manually, since we
                # have a real numerator
                x, y = z_r[ix,iy], z_i[ix,iy]
                x2, y2 = x*x, y*y
                p2 = x2 - y2
                f = 1/(p2*p2 + 4*x2*y2)
                z2inv_r, z2inv_i = f*p2, -2*f*x*y  # 1/z^2
                x, y = one_third*(2*x+z2inv_r), one_third*(2*y+z2inv_i)
                z_r[ix,iy], z_i[ix,iy] = x, y
                s[ix,iy] += sf*abs(x*x+y*y-1)
            end
        end
    end
    nothing
end

alloc{T}(::Type{T}, m, n=m) = zeros(T, m, n), zeros(T, m, n), zeros(T, m, n), Array(RGB{N0f8}, m, n)

function colorize!(col::Matrix{RGB{N0f8}}, s, z_r, z_i)
    for i in eachindex(s)
        if s[i] > sthresh
            col[i] = RGB(1,1,1)
        else
            sm = smap(s[i])
            if z_r[i] > 0
                col[i] = HSV(0,sm,1)
            else
                if z_i[i] > 0
                    col[i] = HSV(120,sm,1)
                else
                    col[i] = HSV(240,sm,1)
                end
            end
        end
    end
    col
end

smap(s) = clamp(log(sthresh/s)/5, 0, 1)


using FFTW
using LinearAlgebra
using DSP
using Plots
using Statistics
using ProgressMeter
include("pcfm.jl")

function autocorr(x)
  return conv(x, conj.(x[end:-1:1]))
end

function ∇J(B, x, u, l)
  """
  ∇J(B,Bb,x,m,u,a)
  Compute the frequency template error (FTE) and its associated gradient

  # Arguments
  - `B::Array`: M X M Orthogonal basis
  - `x::Vector`: M X 1 phase code vector
  - `u::Vector`: 2M-1 X 1 Frequency template
  - `l::Integer`: Norm to use for error calculation
  """
  # Oversampled phase code length
  m = size(B, 1)
  # Zero-pad basis matrix to length 2M-1 (for FFT)
  Bb = vcat(B, zeros(m - 1, trunc(Int, size(B, 2))))
  # PCFM representation of the input phase code vector
  s = exp.(im .* B * x)
  # Pad the waveform to length 2M-1
  sb = vcat(s, zeros(m - 1, size(s, 2)))
  # Compute the (normalized) PSD of the PCFM waveform
  sbf = fftshift(fft(sb))
  sbf = sbf ./ maximum(abs.(sbf))
  # FTE calculation
  J = norm(abs.(sbf) .^ 2 .- u, l)
  # Return the error and gradient
  return (J,
    2 / (J) .* transpose(Bb) * imag.(conj.(sb) .* ifft(ifftshift((abs.(sbf) .^ 2 .- u) .* sbf))))
end

function ∇logJ(B, x, u, a, l)
  """
  ∇logJ(B,Bb,x,m,u,a)
  Compute the log frequency template error (log-FTE) and its associated gradient

  # Arguments
  - `B::Array`: M X M Orthogonal basis
  - `x::Vector`: M X 1 phase code vector
  - `u::Vector`: 2M-1 X 1 Frequency template
  - `a::Integer`: Log base for error computation
  - `l::Integer`: Norm to use for error calculation

  """
  # Oversampled phase code length
  m = size(B, 1)
  # Zero-pad basis matrix to length 2M-1 (for FFT)
  Bb = vcat(B, zeros(m - 1, size(B, 2)))
  # PCFM representation of the input phase code vector
  s = exp.(im .* B * x)
  # Pad the waveform to length 2M-1
  sb = vcat(s, zeros(m - 1, size(s, 2)))
  # Compute the (normalized) PSD of the PCFM waveform
  sbf = fftshift(fft(sb))
  sbf = sbf ./ maximum(abs.(sbf), dims = 1)
  # log-FTE calculation
  J = norm(log.(a, abs.(sbf) .^ 2) .- log.(a, u), l)
  # Return the error and gradient
  return (J,
    (2 / (log(a) * J)) .* transpose(Bb) * imag.(conj.(sb) .* ifft(ifftshift((log.(a, abs.(sbf) .^ 2) .- log.(a, u)) .* sbf))))
end

function profm(u, iter)
  """
  profm(u,iter)

  Iteratively optimize the PSD using alternating projections as described in the
  PRO-FM paper.
  """
  pk = exp.(im .* angle.(ifft(ifftshift(u))))
  for ii = 1:iter
    rk = ifft(ifftshift(abs.(u) .* exp.(im .* angle.(fftshift(fft(pk))))))
    pk = exp.(im .* angle.(rk))
  end
  return pk
end

function optimize(u, nWaveforms, k; a = 10, tol = 1e-5, maxIter = 1000, savePlots = false, filename = "anim_fps60.gif")
  """
  optimize(u,a,tol,maxIter)

  # Arguments
  - `u::Matrix`: 2M-1 X 1 Frequency template
  - `a::Int`: Log base for log-FTE (ignored if regular FTE is used)
  - `tol::Float64`: Tolerance for early stopping
  - `maxIter::Integer`: Maximum number of iterations
  """
  # TODO: Allow the user to decide which optimization method to use

  #Calculate m from u.
  m = trunc(Int, (length(u) + 1) / 2)
  # Get a randomly initialized phase change vector and phase shaping basis
  # functions from PCFM generator
  (_, x, B) = pcfm(m, k, nWaveforms)
  # Gradient descent Parameters
  μ = 0.5
  β = 0.9
  # Store the error at each iteration
  Jvec = ones(maxIter - 1, 1)
  pkOld = 0
  sbf = zeros(length(u), nWaveforms)
  anim = @animate for ii = 1:maxIter
    # Heavy-ball gradient descent
    (J, ∇) = ∇J(B, x, u, 2)
    # (J, ∇) = ∇logJ(B, x, u, a, 2)
    Jvec[ii:end] .= J
    # Update the search direction
    if ii == 1
      pk = ∇
    else
      pk = ∇ .+ β .* pkOld
    end
    # Update x
    x -= μ .* pk
    # Stopping condition
    if all(abs.(pk .- pkOld) .< tol)
      break
    end
    pkOld = pk
    # Plots
    # Compute and plot the PSD
    s = exp.(im .* B * x)
    sb = vcat(s, zeros(m - 1, nWaveforms))
    sbf = fftshift(fft(sb))
    sbf = mean(sbf, dims = 2)
    sbf = sbf ./ maximum(abs.(sbf), dims = 1)
    p1 = plot(10 * log10.(abs.(sbf) .^ 2), ylim = (-50, 0),
      xlabel = "Sample Index", label = "Actual PSD")
    plot!(10 * log10.(u), ylim = (-50, 0), label = "PSD Template")
    # Compute and plot the autocorrelation
    corr = abs.(autocorr(s)) ./ maximum(abs.(autocorr(s)))
    p2 = plot(10 * log10.(corr), ylim = (-30, 0), xlabel = "Sample Index",
      ylabel = "Magnitude (dB)", label = "Autocorrelation")
    # Compute and plot the current error
    p3 = plot(Jvec, ylim = (0, maximum(Jvec)), xlabel = "Iteration Number", ylabel = "Error", label="Error")
    plot(p1, p2, p3, layout = (3, 1), linewidth = 2)
  end
  if savePlots
    gif(anim, filename, fps = 60)
  end
  s = exp.(im .* B * x)
  return (x, s)

end


###
# Specialized functions for Newsvendor with Ber(p) demand
# unit holding, backorder cost: b = s/(1-s)
# Because of Bernoulli demand, represent data as mhat = count zeta ==1
# xk = 1 iff p(alpha) > 1-s

# prob xk = 1
nv_prob_k(p, p0, alpha, Nhat_k, s) = ccdf(Binomial(Nhat_k, p), (Nhat_k + alpha) * (1 - s) - alpha * p0)

#vectorized way to simulate counts
#returns vector mhat_k for 1s only.
nv_sim_path(ps, Ns) = @.sum(rand(Bernoulli(ps), Ns))

#E[Z_k(alpha)], i.e. expected true obj at x(alpha)
function exp_nv_obj_k(p, p0, alpha, Nhat_k, s)
    b = s / (1 - s)
    (1 - p - b * p) * nv_prob_k(p, p0, alpha, Nhat_k, s) + b * p
end

#expected value of fullInfo solution
function exp_nv_fullInfo_k(p, s)
	b = s / (1 - s)
	xstar = p > 1 - s ? 1. : 0.
	b * p + (1 - p - b * p) * xstar
end

#Z_k(alpha)..  mhat is # zeta-k == 1
function nv_obj_k(mhat, p, p0, alpha, Nhat_k, s)
	p_alpha = shrink(mhat / Nhat_k, p0, alpha, Nhat_k)
	if p_alpha > 1 - s #xk = 1
		return 1 - p
	else #xk = 0
		return s / (1 - s) * p
	end
	return 0.  #never reached
end

#Computes Nint * ZLoo_k(alpha)
#Nint is the overall intensity for observations
#Assumes lambdas = 1
function nv_loo_k(mhat_k, p0, alpha, Nhat_k, s)
	out = 0.
	#i corresponding to demand = 1
	if mhat_k > 0
		p_alpha = shrink((mhat_k - 1)/(Nhat_k - 1), p0, alpha, Nhat_k - 1)
		#only non-zero if x neq 1
		out += p_alpha > 1 - s ? 0. : mhat_k * s / (1 - s)
	end

	#i corresponding to demand = 0
	if mhat_k < Nhat_k
		p_alpha = shrink(mhat_k / (Nhat_k - 1), p0, alpha, Nhat_k - 1)
		#only non-zero if x  = 1
		out += p_alpha > 1 - s ? (Nhat_k - mhat_k) : 0.
	end
	return out
end

#this is a funny saa-ish term in the Loo-Stability decomposition
#Returns Nint * saa-ish term.
function nv_saa_k(mhat_k, p0, alpha, Nhat_k, s)
	p_alpha = shrink(mhat_k/Nhat_k, p0, alpha, Nhat_k)
	if p_alpha > 1 - s #xk = 1
		return Nhat_k - mhat_k
	else #xk = 0
		return s / (1 - s) * mhat_k  #Note, not scaling by Nhat_k!
	end
	return 0.  #never reached
end

#Optimzes Z(\alpha) over alpha_grid
#returns alphaOR, minimizingAlphaIndex, curveInAlpha
function nv_oracle_alpha(mhats, ps, p0, alpha_grid, Nhats, ss)
	out = map(a-> nv_obj(mhats, ps, p0, a, Nhats, ss), alpha_grid)
	jstar = argmin(out)
	return alpha_grid[jstar], jstar, out
end

#returns p0OR, alphaOR, jOR for alpha
function nv_oracle_both(mhats, ps, p0_grid, alpha_grid, Nhats, ss)
	out = zeros(length(alpha_grid))
	or_perf, p0OR, alphaOR, jOR = Inf, 0., 0., 0
	for p0 in p0_grid
		out[:] = map(a-> nv_obj(mhats, ps, p0, a, Nhats, ss), alpha_grid)
		jstar = argmin(out)
		if out[jstar] < or_perf
			jOR = jstar
			or_perf = out[jstar]
			p0OR = p0
			alphaOR = alpha_grid[jstar]
		end
	end
	return p0OR, alphaOR, jOR
end

#optimizes ZLoo(\alpha)over alpha_grid
#Note curve in alpha = Nintensity * lamavg * ZLoo(alpha)
#returns alpha^LOO, maximizingAlphaIndex, curveInAlpha
function nv_loo_alpha(mhats, p0, alpha_grid, Nhats, ss)
	out = map(a-> nv_loo(mhats, p0, a, Nhats, ss), alpha_grid)
	jstar = argmin(out)
	return alpha_grid[jstar], jstar, out
end

#loo estimate optimizing boht p0 and alpha0
function nv_loo_both(mhats, p0_grid, alpha_grid, Nhats, ss)
	out = zeros(length(alpha_grid))
	loo_perf = Inf
	p0LOO = 0.
	alphaLOO = 0.
	jLOO = -1
	for p0 in p0_grid
		out[:] = map(a-> nv_loo(mhats, p0, a, Nhats, ss), alpha_grid)
		jstar = argmin(out)
		if out[jstar] < loo_perf
			jLOO = jstar
			loo_perf = out[jstar]
			p0LOO = p0
			alphaLOO = alpha_grid[jstar]
		end
	end
	return p0LOO, alphaLOO, jLOO
end

#Optimzes E[Z(\alpha)] over alpha_grid
#returns alphaAP, minimizingAlphaIndex, curveInAlpha
function nv_apriori_alpha(ps, p0, alpha_grid, Nhats, ss)
	out = map(a-> exp_nv_obj(ps, p0, a, Nhats, ss), alpha_grid)
	jstar = argmin(out)
	return alpha_grid[jstar], jstar, out
end

#vector versions assume lam_k = 1 for all k
nv_loo(mhats, p0, alpha, Nhats, ss) =  mean( nv_loo_k.(mhats, p0, alpha, Nhats, ss) )
nv_obj(mhats, ps, p0, alpha, Nhats, ss) = mean( nv_obj_k.(mhats, ps, p0, alpha, Nhats, ss) )
nv_saa(mhats, p0, alpha, Nhats, ss) = mean( nv_saa_k.(mhats, p0, alpha, Nhats, ss) )

exp_nv_obj(ps, p0, alpha, Nhats, ss) = mean( exp_nv_obj_k.(ps, p0, alpha, Nhats, ss) )
nv_var_obj(ps, p0, alpha, Nhats, ss) = mean( nv_var_obj_k.(ps, p0, alpha, Nhats, ss) )
exp_nv_fullInfo(ps, s) = mean(exp_nv_fullInfo_k.(ps, s))


export tridiag_real
using LinearAlgebra

"""
    tridiag_real(A::AbstractMatrix, d::Integer)

The symmetric real or Hermitian input matrix `A` is transformed to a real
`SymTridiagonal` matrix.
For general matrices only `d` superdiagonals are processed.
Works efficiently for `BandedMatrices`.
"""
function tridiag_real(A::AbstractMatrix, d::Integer)
    m, n = size(A)
    m == n || throw(DimensionMismatch("matrix is not square: dimensions are $((m, n))"))
    0 < d <= n || throw(ArgumentError("number of diagonals $d not in 1:$n"))

    B = copybands(A, d)
    td = make_tridiag!(B)
    dv = [real(B[i,1]) for i in 1:n]
    ev = d > 1 ? [abs(B[i,2]) for i = 2:n] : zeros(real(eltype(A)), n-1)
    SymTridiagonal(dv, ev)
end

# B[1:q,1:n] has the bands of a (2q-1)-banded hermitian nxn-matrix.
# superdiagonal i valid in [i+1:n].
function make_tridiag!(B)
    n, d = size(B)
    0 < d <= n || throw(ArgumentError("number of diagonals $d not in 1:$n"))

    @inbounds for bm = d-1:-1:2
        for k = 1:n-bm
            kp = k
            apiv = B[bm+kp,bm+1]
            iszero(apiv) && continue
            for i = bm+k-1:bm:n-1
                b = B[i, i-kp+1]
                c, s, r, α = givens2(b, apiv)
                u, v = B[i,1], B[i+1,1]
                upx = (u + v) / 2
                B[i,1] = (u - v) / 2
                B[i,i-kp+1] = r
                for j = kp+1:i
                    u = B[i,i-j+1]
                    v = B[i+1,i-j+2] * α'
                    B[i,i-j+1], B[i+1,i-j+2] = u * c + v * s, -u * s + v * c
                end
                B[i+1,1] = -(B[i,1] * α)'
                ip = i + bm
                for j = i+1:min(ip, n)
                    u = B[j,j-i+1]
                    v = B[j,j-i] * α
                    B[j,j-i+1], B[j,j-i] = u * c + v * s, -u * s + v * c
                end
                w = real(B[i+1,1])
                B[i,1] = upx - w
                B[i+1,1] = upx + w
                if ip < n
                    v = B[ip+1,ip-i+1] * α
                    apiv, B[ip+1,ip-i+1] = v * s, v * c
                end
                kp = i
            end
        end
    end
    B
end

function givens2(a::T, b::T) where T <: AbstractFloat
    r = _hypot(a, b)
    s = b / r
    c = a / r
    c, s, r, one(T)
end
@inline function givens2(a::T, b::T) where T <: Complex
    aa = _abs(a)
    ba = _abs(b)
    ra = _hypot(aa, ba)
    s = ba / ra
    c = aa / ra
    ua = a / aa
    α = ua' * (b / ba)
    r = ua * ra
    c, s, r, α
end

# speedy versions of abs and hypot (in the common case)
function _abs(x::Complex)
    s = sqrt(abs2(real(x)) + abs2(imag(x)))
    _isclean(s) ? s : abs(x)
end
_abs(x) = abs(x)
function _hypot(x, y)
    s = sqrt(abs2(x) + abs2(y))
    _isclean(s) ? s : hypot(x, y)
end

function _isclean(a::T) where T<:AbstractFloat
    isfinite(a) &&
    a >= sqrt(2 * floatmin(T))
end

copybands(A, d) = copybands!(zeros(eltype(A), size(A,1), d), A)

function copybands!(B::AbstractMatrix, A::AbstractMatrix)
    n, d = size(B)
    for i = 1:d
        for j = i:n
            B[j,i] = A[j-i+1,j]
        end
    end
    B
end


using Crux
using Test
using Zygote
using Flux

## gradient penalty
m = Dense(2,1, init=ones, bias=false)
x = ones(Float32, 2, 100)
@test gradient_penalty(m, x) ≈ (sqrt(2) - 1)^2

# Other gradient penalty test

idim = 5
batch_size = 8
m = Chain(Dense(idim, 2*idim, tanh), Dense(2*idim,1)) |> gpu
x = rand(Float32, idim, batch_size) |> gpu
y = rand(Float32, 1, batch_size) |> gpu

function total_loss()
    Flux.mse(m(x), y) + gradient_penalty(m, x)
end

l, b = Flux.pullback(total_loss, Flux.params(m))
grad = b(1f0)
grad.grads


using DynamicHMC, LogDensityProblems
using Distributions, Parameters, Random
using StatsPlots, PositiveFactorizations, Calculus
using Flux

include(string(pwd(),"/src/simulation.jl"))
include(string(pwd(),"/gali_bayesian.jl"))
include(string(pwd(),"/mcmc/priors_v2.jl"))
include(string(pwd(),"/mcmc/foos.jl"))

yy,shocks = simulate_dsge(GAMMA_0,GAMMA_1,PSI,PI,500)

prob = DSGE_Model(yy[:,2],1/3)

prob((bet = 0.99,epsilon = 6,theta=2/3,sig=1,s2=1,phi=1,phi_pi=1.5,phi_y=0.5/4,rho_v=0.5))

t = problem_transform(prob)

P = TransformedLogDensity(t,prob)

aa = [0.99,6,2/3,1,1,1,1.5,0.5/4,0.5]

true_vals = TransformVariables.inverse(t,(bet = 0.99,epsilon = 6,theta=2/3,sig=1,s2=1,phi=1,phi_pi=1.5,phi_y=0.5/4,rho_v=0.5))

LogDensityProblems.logdensity(P,true_vals)

#grads = Calculus.gradient(P)

function LogDensityProblems.capabilities(::TransformedLogDensity{TransformVariables.TransformTuple{NamedTuple{(:bet, :epsilon, :theta, :sig, :s2, :phi, :phi_pi, :phi_y, :rho_v),Tuple{TransformVariables.ScaledShiftedLogistic{Float64},TransformVariables.ShiftedExp{true,Float64},TransformVariables.ScaledShiftedLogistic{Float64},TransformVariables.ShiftedExp{true,Float64},TransformVariables.ShiftedExp{true,Float64},TransformVariables.ShiftedExp{true,Float64},TransformVariables.ShiftedExp{true,Float64},TransformVariables.ShiftedExp{true,Float64},TransformVariables.ScaledShiftedLogistic{Float64}}}},DSGE_Model{Array{Float64,1},Float64}})
    LogDensityProblems.LogDensityOrder{1}() # can do gradient
end

#LogDensityProblems.dimension(::NormalPosterior) = 2 # for this problem

function LogDensityProblems.logdensity_and_gradient(problem::TransformedLogDensity{TransformVariables.TransformTuple{NamedTuple{(:bet, :epsilon, :theta, :sig, :s2, :phi, :phi_pi, :phi_y, :rho_v),Tuple{TransformVariables.ScaledShiftedLogistic{Float64},TransformVariables.ShiftedExp{true,Float64},TransformVariables.ScaledShiftedLogistic{Float64},TransformVariables.ShiftedExp{true,Float64},TransformVariables.ShiftedExp{true,Float64},TransformVariables.ShiftedExp{true,Float64},TransformVariables.ShiftedExp{true,Float64},TransformVariables.ShiftedExp{true,Float64},TransformVariables.ScaledShiftedLogistic{Float64}}}},DSGE_Model{Array{Float64,1},Float64}}, x)
    logdens = LogDensityProblems.logdensity(P,x)
    grad = -1*Flux.gradient(y->model_flux(y)[1],x)[1] .*exp.(model_flux(x)[1])
    logdens, grad
end

res = @time mcmc_with_warmup(Random.GLOBAL_RNG, P,1000)

posterior = transform.(t, res.chain)

post_array = zeros(1000,9)

for i in 1:1000
    @unpack bet, epsilon, theta, sig, s2,phi,phi_pi,phi_y,rho_v = posterior[i,:][1]
    post_array[i,:] = [bet epsilon theta sig s2 phi phi_pi phi_y rho_v]
end

StatsPlots.density(post_array[:,1], legend = :none)
title!(latexify("beta()"))
vline!([true_pars[:bet]])
vline!([mean(post_array[:,1])])

StatsPlots.density(post_array[:,2], legend = :none)
title!(latexify("epsilon()"))
vline!([true_pars[:epsilon]])
vline!([mean(post_array[:,2])])

StatsPlots.density(post_array[:,3], legend = :none)
title!(latexify("theta()"))
vline!([true_pars[:theta]])
vline!([mean(post_array[:,3])])

StatsPlots.density(post_array[:,4], legend = :none)
title!(latexify("sigma()"))
vline!([true_pars[:sig]])
vline!([mean(post_array[:,4])])

StatsPlots.density(post_array[:,5], legend = :none)
title!(latexify("sigma()^2"))
vline!([true_pars[:s2]])
vline!([mean(post_array[:,5])])

StatsPlots.density(post_array[:,6], legend = :none)
title!(latexify("phi()"))
vline!([true_pars[:phi]])
vline!([mean(post_array[:,6])])

StatsPlots.density(post_array[:,7], legend = :none)
title!(latexify("phi()_pi",cdot = false))
vline!([true_pars[:phi_pi]])
vline!([mean(post_array[:,7])])

StatsPlots.density(post_array[:,8], legend = :none)
title!(latexify("phi()_y",cdot = false))
vline!([true_pars[:phi_y]])
vline!([mean(post_array[:,8])])

StatsPlots.density(post_array[:,9], legend = :none)
title!(latexify("rho()_v",cdot = false))
vline!([true_pars[:rho_v]])
vline!([mean(post_array[:,9])])


using Pkg
Pkg.activate(".")
using Soss

myModel = @model begin
    N ~ Poisson(100)
    K ~ Poisson(2.5)
    p ~ Dirichlet(K, 1.0)
    μ ~ Normal(0,1.5) |> iid(K)
    σ ~ HalfNormal(1)
    θ = [(m,σ) for m in μ]
    x ~ MixtureModel(Normal, θ, p) |> iid(N)
end

using Parameters
using Plots
pyplot()

using Random
using StatsPlots

m = @model begin
    p ~ Uniform()
    μ ~ Normal(0, 1.5) |> iid(2)
    σ ~ HalfNormal(1)
    θ = [(m, σ) for m in μ]
    x ~ MixtureModel(Normal, θ, [p,1-p]) |> iid(100)
end

Random.seed!(20)
grt = rand(m)
mplot(grt)

function mplot(data)
    @unpack μ,σ,p,x = data
    p = [p, 1-p]
    lo,hi = extrema(x)
    lo=min(lo,-4)
    hi=max(hi,4)
    xs = range(lo,hi,length=1000)
    dist = MixtureModel(Normal, tuple.(μ,σ), p)
    plt = plot(xs, p[1] .* pdf.(Normal(μ[1],σ), xs)
        , legend=false, fill = (0, 0.1, :red), linecolor=:red
        , axis=nothing)
    plot!(plt, xs, p[2] .* pdf.(Normal(μ[2],σ), xs)
        , fill = (0, 0.1, :blue), linecolor=:blue)
    plot!(plt,xs, pdf.(dist, xs), legend=false, linewidth=2, linecolor=:black)
    scatter!(plt, x, zeros(size(x))
        , marker=:vline, markercolor=:black, markersize=10)
    xlims!((-4,4))
end

r = makeRand(m)
anim = @animate for i=1:100
    mplot(r())
end
gif(anim, "m.gif", fps = 1)

#############################
# Inverse model (posterior distribution)

m_fwd = m(:p,:μ,:σ)
r = makeRand(m_fwd)
anim = @animate for i=1:100
    v = merge(grt, r(; grt...))
    mplot(v)
end
gif(anim, "m_fwd.gif", fps = 3)

m_inv = m(:x)
post = nuts(m_inv; grt...)
post = post.samples

anim = @animate for par in post[1:100]
    par = merge(grt, par)
    mplot(par)
end
gif(anim, "m_inv.gif", fps = 1)

####################################
# Posterior predictive checks

using StatsBase
xs  = range(-4,4,length=100)
ppc = repeat(grt.x, 1, 1000)
r = makeRand(m_fwd)

plt = plot(xs,ecdf(grt.x).(xs)
    ,linewidth=3,legend=false,size=(1200,800), axis=nothing)
# Each column is a posterior prediction
for j in 1:1000
    ppc = r(;post[j]...).x
    plot!(plt, xs, (rand(100) ./ 100 .+ ecdf(ppc).(xs)) ./ 1.01,linewidth=3, linecolor=:black, linealpha=0.01)
end
plot!(xs,ecdf(grt.x).(xs),linewidth=5,legend=false,linecolor=:white)
plot!(xs,ecdf(grt.x).(xs),linewidth=4,legend=false,linecolor=:orange)
plt
savefig("ppc.png")