#=

  Global constraint global contiguity regular in Julia ConstraintSolver.jl

  From Global Constraint Catalogue
  http://www.emn.fr/x-info/sdemasse/gccat/Cglobal_contiguity.html
  """
  Enforce all variables of the VARIABLES collection to be assigned to 0 or 1.
  In addition, all variables assigned to value 1 appear contiguously.

  Example:
  (<0,1,1,0>)

  The global_contiguity constraint holds since the sequence 0 1 1 0 contains
  no more than one group of contiguous 1.
  """

  Here's a variant of the global constraint global contiguity
  using (a decomposition of) the regular constraint.

  For n=0:13 the number of solutions is
  [1, 2, 4, 7, 11, 16, 22, 29, 37, 46, 56, 67, 79]

  Which is the (or rather one of the possible) OEIS sequence https://oeis.org/A000124
  """
  Central polygonal numbers (the Lazy Caterer's sequence): n(n+1)/2 + 1; or,
  maximal number of pieces formed when slicing a pancake with n cuts.
  """

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

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

function global_contiguity_test(n,print_solutions=true,all_solutions=true)

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

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

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

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

                                                            # https://wikunia.github.io/ConstraintSolver.jl/stable/options/#branch_strategy-(:Auto)
                                                            "branch_strategy" => :IMPS, # default
                                                            # "branch_strategy" => :ABS, # Activity Based Search
                                                            # "activity.decay" => 0.999, # default 0.999
                                                            # "activity.max_probes" => 10, # default, 10
                                                            # "activity.max_confidence_deviation" => 20, # default 20

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

                                                            "time_limit"=>6,

                                                            # "backtrack" => false, # default true
                                                            # "backtrack_sorting" => false, # default true

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

    @variable(model, x[1:n], Bin)

    y = global_contiguity_regular(model, x)

    # Solve the problem
    optimize!(model)

    status = JuMP.termination_status(model)
    # println("status:$status")
    num_sols = 0
    if status == MOI.OPTIMAL
        num_sols = MOI.get(model, MOI.ResultCount())
        println("num_sols:$num_sols\n")
        if print_solutions
            for sol in 1:num_sols
                println("solution #$sol")
                x_val = convert.(Integer,JuMP.value.(x; result=sol))
                y_val = convert.(Integer,JuMP.value.(y; result=sol))
                println("x:$x_val")
                println("y:$y_val")
                println()

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

    return status, num_sols
end

sols = []
for n in 0:13
    global sols
    @time status, num_sols = global_contiguity_test(n,true,true)
    if status == MOI.OPTIMAL
        push!(sols,num_sols)
    end
end
println("\nsols:$sols")

# @time global_contiguity_test(10,true,true)


using Distances, InvertedIndices

export LCS, FastLCS
export print_info, get_info
export LSP, FastLSP
## Interaction Distances

struct LCS <: Metric end
struct FastLCS <: Metric
    curr_row::Vector{Int}
    prev_row::Vector{Int}
    function FastLCS(K::Int)
        new(zeros(Int,K), zeros(Int,K))
    end
end

function Base.show(io::IO, d::FastLCS)
    println(io, "LCS (max interaction len. $(length(d.curr_row)))")
end
struct FastNormLCS <: Metric
    curr_row::Vector{Int}
    prev_row::Vector{Int}
    function FastNormLCS(K::Int)
        new(zeros(Int,K), zeros(Int,K))
    end
end

# LCS
function (dist::LCS)(X::Vector{T},Y::Vector{T})::Float64 where {T}

    n = length(X)
    m = length(Y)

    prevRow = 0:m
    currRow = zeros(Int, m + 1)
    @inbounds begin
        for i = 1:n
            currRow[1] = i
            for j = 1:m
                if X[i] == Y[j]
                    currRow[j+1] = prevRow[j]
                else
                    currRow[j+1] = min(currRow[j], prevRow[j+1]) + 1
                end
            end
            prevRow = copy(currRow)
        end
    end
    return currRow[end]
end

function (dist::LCS)(X::Tuple,Y::Tuple)::Float64

    n = length(X)
    m = length(Y)

    prevRow = 0:m
    currRow = zeros(Int, m + 1)
    @inbounds begin
        for i = 1:n
            currRow[1] = i
            for j = 1:m
                if X[i] == Y[j]
                    currRow[j+1] = prevRow[j]
                else
                    currRow[j+1] = min(currRow[j], prevRow[j+1]) + 1
                end
            end
            prevRow = copy(currRow)
        end
    end
    return currRow[end]
end

# With storage
function (d::FastLCS)(
    X::Vector{T},Y::Vector{T}
    )::Float64 where {T}

    n = length(X)
    m = length(Y)

    @assert (n>0) & (m>0) "both paths must be either of type Nothing or of nonzero length."

    prev_row = view(d.prev_row, 1:(m+1))
    curr_row = view(d.curr_row, 1:(m+1))

    copy!(prev_row, 0:m)
    curr_row .= 0

    @inbounds begin
        for i = 1:n
            curr_row[1] = i
            for j = 1:m
                if X[i] == Y[j]
                    curr_row[j+1] = prev_row[j]
                else
                    curr_row[j+1] = min(curr_row[j], prev_row[j+1]) + 1
                end
            end
            copy!(prev_row, curr_row)
        end
    end
    return curr_row[m+1]
end

# Distances to the null
function (dist::Union{LCS,FastLCS})(X::Nothing, Y::Vector{T})::Float64 where {T}
    return length(Y)
end
function (dist::Union{LCS,FastLCS})(X::Vector{T}, Y::Nothing)::Float64 where {T}
    return length(X)
end
function (dist::Union{LCS,FastLCS})(X::Nothing, Y::Nothing)::Float64 where {T}
    return 0.0
end

# Get locations of longest common subseq (for visuals)

function get_info(
    d::Union{LCS, FastLCS},
    X::Vector{T}, Y::Vector{T}
    ) where {T}

    C = zeros(Int, length(X)+1, length(Y)+1)
    indx, indy = (zeros(Bool, length(X)), zeros(Bool, length(Y)))
    C[:,1] = [i for i = 0:length(X)]
    C[1,:] = [i for i = 0:length(Y)]

    for j = 1:length(Y)
        for i = 1:length(X)
            C[i+1,j+1] = minimum([
            C[i,j] + 2*(X[i] != Y[j]),
            C[i,j+1] + 1,
            C[i+1,j] + 1
            ])
        end
    end

    i = length(X)+1; j = length(Y)+1;
    while (i ≠ 1) & (j ≠ 1)
        if C[i,j] == (C[i-1,j] + 1)
            i = i-1
        elseif C[i,j] == (C[i, j-1] + 1)
            j = j-1
        elseif C[i,j] == (C[i-1, j-1])
            indx[i-1] = true
            indy[j-1] = true
            i = i-1; j = j-1
        end
    end

    return indx, indy
end

# Longest Common Subpath (LSP)

struct LSP <: Metric end

function (dist::LSP)(X::Vector{T},Y::Vector{T})::Float64 where {T}

    # Here we consider a Dynamic programming approach
    n = length(X)
    m = length(Y)

    @assert (n>0) & (m>0) "both paths must be either of type Nothing or of nonzero length."

    prev_row = zeros(Float64, m + 1)
    curr_row = zeros(Float64, m + 1)
    z = 0.0

    @inbounds begin
        for i = 1:n
            for j = 1:m
                if X[i] == Y[j]
                    curr_row[j+1] = prev_row[j] + 1.0 # Subpath length increment
                    if curr_row[j+1] > z
                        z = curr_row[j+1]
                    end
                else
                    curr_row[j+1] = 0.0
                end
            end
            copy!(prev_row, curr_row)
        end
    end
    return n + m - 2*z
end

struct FastLSP <: Metric
    curr_row::Vector{Float64}
    prev_row::Vector{Float64}
    function FastLSP(K::Int)
        new(zeros(Float64,K), zeros(Float64,K))
    end
end

function (dist::FastLSP)(X::Vector{T},Y::Vector{T})::Float64 where {T}
    # Here we take a Dynamic programming approach, but use pre-allocated arrays for storage.
    n = length(X)
    m = length(Y)

    @assert (n>0) & (m>0) "both paths must be either of type Nothing or of nonzero length."

    prev_row = view(dist.prev_row, 1:(m+1))
    curr_row = view(dist.curr_row, 1:(m+1))

    prev_row .= 0.0
    curr_row .= 0.0
    z = 0.0

    @inbounds begin
        for i = 1:n
            for j = 1:m
                if X[i] == Y[j]
                    curr_row[j+1] = prev_row[j] + 1.0 # Subpath length increment
                    if curr_row[j+1] > z
                        z = curr_row[j+1]
                    end
                else
                    curr_row[j+1] = 0.0
                end
            end
            copy!(prev_row, curr_row)
        end
    end
    return n + m - 2*z
end

function (dist::Union{LSP,FastLSP})(X::Nothing, Y::Vector{T})::Float64 where {T}
    return length(Y)
end
function (dist::Union{LSP,FastLSP})(X::Vector{T}, Y::Nothing)::Float64 where {T}
    return length(X)
end
function (dist::Union{LSP,FastLSP})(X::Nothing, Y::Nothing) where {T}
    return 0.0
end

function get_info(
    d::Union{LSP, FastLSP},
    X::Vector{T}, Y::Vector{T}
    ) where {T}

    n,m = (length(X), length(Y))
    curr_row = zeros(Int, m+1)
    prev_row = zeros(Int, m+1)
    z = 0

    ix, iy = (0,0) # Denotes end of common subpath in X and Y resp.

    begin
        for i = 1:n
            for j = 1:m
                if X[i] == Y[j]
                    curr_row[j+1] = prev_row[j] + 1 # Subpath length increment
                    if curr_row[j+1] > z
                        z = curr_row[j+1] # Found new longest subpath
                        ix, iy = (i,j)
                    end
                else
                    curr_row[j+1] = 0
                end
            end
            copy!(prev_row, curr_row)
        end
    end

    indx = [(ix-z+1 ≤ i ≤ ix) ? true : false for i in 1:n]
    indy = [(iy-z+1 ≤ i ≤ iy) ? true : false for i in 1:m]
    return indx, indy
end

function print_info(
    d::Union{LSP,FastLSP,LCS,FastLCS},
    x::Vector{T}, y::Vector{T}
    ) where {T}

    indx, indy = get_info(d, x, y)
    ix , iy = (0,0)
    while true
        ixerr = isnothing(findnext(indx, ix+1)) ? length(x) : findnext(indx, ix+1)-1
        iyerr = isnothing(findnext(indy, iy+1)) ? length(y) : findnext(indy, iy+1)-1
        for j in (ix+1):ixerr
            println("$(x[j]) (delete)")
        end
        for j in (iy+1):iyerr
            println("$(y[j]) (add)")
        end
        if isnothing(findnext(indx, ix+1))
            break
        end
        ix = findnext(indx, ix+1)
        iy = findnext(indy, iy+1)
        println("$(x[ix]) (keep)")
    end
end

struct SuffixLSP <: Metric end


module Optimization

using Logging
using KrylovMethods

using DeformableRegistration
include("helpers/objectiveFunctionCreation.jl")

export checkStoppingCriteria, ArmijoLineSearch, optimizeGaussNewton

function optimizeGaussNewton(Jfunc::Function,  # objective Function
                             y::Array{Float64,1}, yInitial::Array{Float64,1}, options::regOptions)

    # Initilization of JRef: reference value of objective Function
    JRef = Jfunc(yInitial)[1]
    JOld = NaN

    output = (Logging.LogLevel == Logging.DEBUG) |  (Logging.LogLevel == Logging.INFO)

    for iter = 1:options.maxIterGaussNewton

        # get derivatives of objective function
        J, dJ, d2J = Jfunc(y,doDerivative=true,doHessian=true)

        # solve system of linear equations d2J*dy=-dJ
          # (dy: change of the variables / search direction)
        # number of parameters < 10 => use backslash operator
        # else (large systems) => use conjugate gradient (cg) method

        cgIterations=0

        if(length(y)<10 && !(typeof(d2J) <: Function))
            dy=d2J\-dJ
        else
            dy,flag,resvec,cgIterations = KrylovMethods.cg(d2J,-dJ,maxIter=options.maxIterCG, tol=1e-5)[1:4]
        end

		    # check descent direction
		    if( (dJ'*dy)[1] > 0)
            dy = -dy
            warn("Changing sign of computed descent direction. This is a suspicious move.")
		    end

        # armijo line search (LS) method
        stepLength,LSiter,LSfailed = ArmijoLineSearch(Jfunc,J,dJ,y,dy)
        if(LSfailed)
            break
        end

        # output
        if(cgIterations==0)
          s = @sprintf("%3d: J %8.4e     LSiter: %2d    J/Jref: %1.2f \n",iter, J[1], LSiter, J[1]/JRef[1])
          info(s)
        else
          s = @sprintf("%3d: J %8.4e     LSiter: %2d     CGiter: %3d     J/Jref: %1.2f \n",iter, J[1], LSiter,cgIterations, J[1]/JRef[1])
          info(s)
        end

        # update parameter y
        y = y + stepLength.*dy

        # stopping criteria
        if (iter>1)
            if checkStoppingCriteria(J, JOld, JRef, dJ, y, stepLength*dy) #printActiveStoppingCirteria=output
                break
            end
        end

        # updating JOld
        JOld = J[1]

    end

    return y

end

function ArmijoLineSearch(Jfunc::Function,         # objective function
                          Jc::Float64,             # Jc = Jfunc(y)
                          dJ::Array{Float64,1},    # gradient of the objective funtion
                          y::Array{Float64,1},     # old variables
                          dy::Array{Float64,1};    # change of the variables / search direction
                          tolLS::Float64 = 1e-4)   # tolerance

    stepLength = 1.0; LSiter = 1; LSfailed = false;

    while( (Jfunc(y+stepLength*dy)[1]) >= (Jc + (stepLength.*tolLS.*dJ'*dy)[1]) )

        if(stepLength < 1e-5)
            LSfailed = true
            break
        end

        stepLength = 0.5 * stepLength
        LSiter = LSiter + 1

    end

    if(LSfailed)
        s = @sprintf("   0. Line search failed after %2d iterations.",LSiter)
        Logging.debug(s)
    end

    return stepLength,LSiter,LSfailed

end

function checkStoppingCriteria(J,JOld,JRef,    # value of the current objective function J(y+dy), the old J(y) and the reference J(yInital)
                               dJ,             # gradient of the objective funtion
                               y,              # old variables
                               dy;             # change of the variables / search direction
                               tolJ = 1e-3,    # tolerance: change of the objective function
                               tolY = 1e-2,    # tolerance: change of the variables
                               tolG = 1e-2,    # tolerance: change of the gradient
                               tolQ = 1e-5     # tolerance: change of quotient
                               )

    # Initialize STOP-Array with false
    STOP = Array{Bool}(5); STOP[:] = false

    # Either STOP[1:3] or STOP[4]
    STOP[1] = abs(JOld-J) <= tolJ  * (1+abs(JOld))
    STOP[2] = norm(dy)    <= tolY  * (1+norm(y))
    STOP[3] = norm(dJ)    <= tolG  * (1+abs(JOld))
    STOP[4] = norm(dJ)    <= 1e6 * eps()
    STOP[5] = abs((JOld-J)/(JRef-J)) <= tolQ

    if(all(STOP[1:3]))
        s = @sprintf("STOPPING CRITERIA:\n")
        Logging.info(s)
        s = @sprintf("   1. |JOld-J| = %5e <= %5e \n",abs(JOld-J), tolJ*(1+abs(JOld)) )
        Logging.info(s)
        s = @sprintf("   2. ||dy|| = %5e <= %5e \n",norm(dy),    tolY*(1+norm(y))  )
        Logging.info(s)
        s = @sprintf("   3. ||dJ|| = %5e <= %5e \n",norm(dJ),    tolG*(1+abs(JOld)) )
        Logging.info(s)
    end

    if(STOP[4])
        s = @sprintf("STOPPING CRITERION:\n")
        Logging.info(s)
        s = @sprintf("   4. ||dJ|| = %5e <= %5e \n",norm(dJ), 1e6*eps())
        Logging.info(s)
    end

    if(STOP[5])
        s = @sprintf("STOPPING CRITERION:\n")
        Logging.info(s)
        s = @sprintf("   5. |(JOld-J)/(JRef-J)|= %5e <= %5e \n",abs((JOld-J)/(JRef-J)), tolQ)
        Logging.info(s)
    end

    return all(STOP[1:3]) | STOP[4] | STOP[5]

end

end


using Oscar
using Polymake
using InvertedIndices
using Combinatorics
using LinearAlgebra
# StackyFan struct
export StackyFan, makeStackyFan, addStackStructure
export stackyWeights, getRayStack, encode
# Fan
export toric_blowup, makeSimplicial, makeSmooth, starSubdivision
export getConeRank, getDimension, getConeFaces
export coneMultiplicity, getMultiplicity
export findFaceContainingRay, findMinimalCone, interiorPoints, findBarycenter
export rootConstruction, rootConstructionDistinguished, rootConstructionDistinguishedIndices
# Stacky fan
export stackyBlowup, coneRayDecomposition
export findStackyBarycenter, findStackyRayMatrix
# Utility
export slicematrix, rowMinors, convertIncidenceMatrix, convertToIncidence
export coneConvert, getCones, getConesPolymake
# Display
export plot3dSimpCone, plot2dCone, showSimpFan, showSimpStackyFan, coneVectorOrder, plot3dCone
export showFan, showStackyFan

"""
    Structure to store information of a stacky fan - this is a fan together with a dictionary assigning stacky values to each ray.

# Properties:
-> `fan` - the underlying fan, as a polymake object
-> `scalars` - the array of stacky values
-> `stacks` - a dictionary assigning stacky values to each ray.

"""
struct StackyFan
    fan::Polymake.BigObjectAllocated
    stacks::Dict{String, Int64}
    # Constructors for the StackyFan object
    StackyFan(
        fan::Polymake.BigObjectAllocated,
        stacks::Dict{String, Int64}) = new(fan, stacks)
    StackyFan(
        rays::Array{Int64, 2},
        cones::Array{Array{Int64, 1}, 1},
        scalars::Array{Int64, 1}) = makeStackyFan(rays, cones, scalars)
    StackyFan(
        fan::Polymake.BigObjectAllocated,
        scalars::Array{Int64, 1}) = addStackStructure(fan, scalars)
end

## Helper functions

"""
    makeStackyFan(::Array{Int64,2},::Array{Array{Int64,1},1},::Array{Int64,1}))

Function to generate a stacky fan from a matrix representing rays as row vectors, a vector of vectors representing the rays contained in each cone, and a vector of stacky values to be assigned the rays. The second input should be zero-indexed.

# Examples
```jldoctest
julia> makeStackyFan([1 0; 1 1; 1 2],[[0,1],[1,2]],[2,2,2])
[ 2 ,  2 ,  2 ]
```
"""
function makeStackyFan(
    rays::Array{<:Number,2},
    cones::Array{Array{Int64,1},1},
    scalars::Array{Int64,1})

    # Construct a normal fan from the given rays and cones
    fan = fulton.NormalToricVariety(INPUT_RAYS=rays, INPUT_CONES=cones)

    # Construct the dictionary
    stack_rays = mapslices(encode, convert(Array{Int64,2},Array(Polymake.common.primitive(fan.RAYS))), dims=2)
    pairs = map((x,y) -> (x,y), stack_rays, scalars)
    stacks = Dict(pairs)

    return(StackyFan(fan, stacks))
end

"""
    addStackStructure(::Polymake.BigObjectAllocated, ::Array{Int64, 1})

Function to generate a stacky fan from a given fan and a set of scalars.

# Examples
```jldoctest
julia> X=Polymake.fulton.NormalToricVariety(INPUT_RAYS=[1 0; 1 1; 1 2],INPUT_CONES=[[0,1],[1,2]]);

julia> stackyWeights(addStackStructure(X,[2,2,2]))
[ 2 ,  2 ,  2 ]
```
"""
function addStackStructure(
    fan::Polymake.BigObjectAllocated,
    scalars::Array{Int64, 1})

    # Construct the dictionary
    stack_rays = mapslices(encode, Polymake.common.primitive(fan.RAYS), dims=2)
    pairs = map((x,y) -> (x,y), stack_rays, scalars)
    stacks = Dict(pairs)

    return(StackyFan(fan, stacks))
end

"""
    encode(::Union{Polymake.VectorAllocated{Polymake.Rational},Polymake.VectorAllocated{Polymake.Integer},Vector{Int64}})

    Internal function that converts a vector, representing a ray in the fan,
to a string in order to allow for hashing for the dictionary.

#Examples
```jldoctest
julia> encode([1,0,2,5])
"1,0,2,5"
```
"""
function encode(objects::Union{Polymake.VectorAllocated{Polymake.Rational},Polymake.VectorAllocated{Polymake.Integer},Vector{Int64}})
    return(foldl((x,y) -> string(x, ',', y), objects))
end

"""
    stackyWeights(::StackyFan)

    Returns a list of the stacky weights of the rays of the given stacky fan with the same order as the rays of the fan.

#Examples
```jldoctest
julia> F=makeStackyFan([1 0; 1 1; 1 2; 1 3],[[0,1],[1,2],[2,3]],[1,2,3,4]);

julia> stackyWeights(F)
[ 1 ,  2 ,  3 ,  4 ]
```
"""
function stackyWeights(sf::StackyFan)
    rayMatrix=convert(Array{Int64,2},Array(Polymake.common.primitive(sf.fan.RAYS)))
    #println("Stacky weights ray matrix: $rayMatrix")
    rayList=slicematrix(rayMatrix)
    out=Int64[]
    for ray in rayList
        stack=sf.stacks[encode(ray)]
        push!(out,stack)
    end
    return out
end

## API functions

"""
    getRayStack(::StackyFan, ::Array{Int64, 1})

    Get the scalar associated with a ray in the given stacky fan structure.

# Examples
```jldoctest
julia> F=makeStackyFan([1 0; 1 1; 1 2; 1 3],[[0,1],[1,2],[2,3]],[1,2,3,4]);

julia> getRayStack(F,[1,2])
3
```
"""
function getRayStack(sf::StackyFan, ray::Array{Int64, 1})
    return sf.stacks[encode(ray)]
end

"""
    rootConstruction(::StackyFan, ::Array{Int64, 1})

Given a fan and a set of scalars corresponding to the rays of the fan,
performs a root construction on the fan by multiplying the stack scalars
by the given values.

rootConstruction returns a new StackyFan object, and does not modify the input.

# Examples
```jldoctest

julia> X=Polymake.fulton.NormalToricVariety(INPUT_RAYS=[1 0 0;1 1 0;1 0 1;1 1 1],INPUT_CONES=[[0,1,2],[1,2,3]]);

julia> SX = StackyFan(X, [2,3,5,7]);

julia> stackyWeights(rootConstruction(SX, [1, 4, 2, 1]))
[ 2 ,  12 ,  10 ,  7 ]
```
"""
function rootConstruction(
    sf::StackyFan,
    scalars::Array{Int64, 1})

    # Multiply the scalars of the fan by the given values
    return StackyFan(sf.fan, stackyWeights(sf) .* scalars)
end

"""
    rootConstructionDistinguishedIndices(::StackyFan, ::Array{Int64, 1}, ::Array{Int64, 1})

    Given a fan, the indices of the distinguished rays in the fan rays (as an incidence matrix), and
a set of scalars corresponding to the rays of the fan, performs a root
construction on the fan by multiplying the stack scalars by the given values.

    rootConstructionDistinguishedIndices returns a new StackyFan object,
and does not modify the input.

# Examples
```jldoctest
julia> X=Polymake.fulton.NormalToricVariety(INPUT_RAYS=[1 0 0;1 1 0;1 0 1;1 1 1],INPUT_CONES=[[0,1,2],[1,2,3]]);

julia> SX = StackyFan(X, [2,3,5,7]);

julia> stackyWeights(rootConstructionDistinguishedIndices(SX, [0,1,1,0], [4, 2, 1, 3]))
[ 2 ,  6 ,  5 ,  7 ]
```
"""
function rootConstructionDistinguishedIndices(
    sf::StackyFan,
    distIndices::Array{Int64, 1},
    scalars::Array{Int64, 1})

    numRays = size(sf.fan.RAYS, 1)
    fullScalars = fill(1, numRays)
    for i in 1:numRays
        if distIndices[i]==1 && scalars[i] != 0
            fullScalars[i] = scalars[i]
        end
    end
    # Multiply the scalars of the fan by the given values
    return rootConstruction(sf, fullScalars)
end

"""
    rootConstructionDistinguished(
        ::StackyFan,
        ::Polymake.Matrix{Polymake.Rational},
        ::Array{Int64, 1})

    Given a fan, a set of distinguished rays, and a set of scalars of equal size,
performs a root construction on the fan on the distinguished rays by multiplying
the stack scalars by the given values.

    rootConstructionDistinguished returns a new StackyFan object,
and does not modify the input.

# Examples
```jldoctest
julia> X=Polymake.fulton.NormalToricVariety(INPUT_RAYS=[1 0 0;1 1 0;1 0 1;1 1 1],INPUT_CONES=[[0,1,2],[1,2,3]]);

julia> SX = StackyFan(X, [2,3,5,7]);

julia> distinguished = X.RAYS[[2,3],:];

julia> stackyWeights(rootConstructionDistinguished(SX, distinguished, [4, 2]))
[ 2 ,  12 ,  10 ,  7 ]
```
"""
function rootConstructionDistinguished(
    sf::StackyFan,
    rays::Polymake.Matrix{Polymake.Rational},
    scalars::Array{Int64, 1})

    # Check that the rays and scalars are the same size
    #if (size(rays, 1) != length(scalars))
    #    error("Inputs are not of equal size")
    #end

    encoded_rays = mapslices(encode, rays, dims=2)
    # Make a copy of the dictionary
    newStacks = copy(sf.stacks)
    for i in 1:length(encoded_rays)
        ray = encoded_rays[i]
        # Multiply the scalar of the corresponding ray
        newStacks[ray] *= scalars[i]
    end

    # Convert the dictionary to an array of scalars matching the indices
    #newScalars = mapslices(ray -> newStacks[encode(ray)], sf.fan.RAYS, dims=2)
    newScalars = Array{Int64, 1}()
    for i in 1:size(sf.fan.RAYS, 1)
        push!(newScalars, newStacks[encode(sf.fan.RAYS[i,:])])
    end

    return StackyFan(sf.fan, newScalars)
end

"""
    findBarycenter(::Union{AbstractSet,AbstractVector},::Polymake.BigObjectAllocated)

    Takes a normal toric variety X and a set s corresponding to a subset of rays of X, and outputs a polymake vector corresponding to the barycenter of those rays.

# Examples
```jldoctest makeSmoothWithDependencies
julia> X=Polymake.fulton.NormalToricVariety(INPUT_RAYS=[1 0;1 1; 0 1],INPUT_CONES=[[0,1],[1,2]]);

julia> s=[1,2];

julia> findBarycenter(s,X)
pm::Matrix<pm::Integer>
2 1
```
"""
function findBarycenter(s::Union{AbstractSet,AbstractVector},X::Polymake.BigObjectAllocated)
#     rayMatrix=convert(Array{Int64,2},Array(Polymake.common.primitive(X.RAYS)))
#     rays = rowMinors(rayMatrix, s)
#     dim=size(rays,2)
#     bary=zeros(Int64,dim,1)
#     for i in 1:size(rays,1)
#         bary+=rays[i,:]
#     end
#     bary=Polymake.common.primitive(bary)
#     return vec(bary)
# end

    rays = rowMinors(Array(X.RAYS), s)
    dim=size(rays,2)
    bary=zeros(Polymake.Rational,dim,1)
    for i in 1:size(rays,1)
        bary+=rays[i,:]
    end
    bary=Polymake.common.primitive(transpose(bary))
    return bary
end

"""
    findStackyBarycenter(::Union{AbstractSet,AbstractVector},::StackyFan)

    Takes a stacky fan SX and a set s corresponding to a subset of rays of SX, calculates the 'stacky rays' corresponding to those rays (the rays times their stacky values), and find the barycenter of the stacky rays.

# Examples
```jldoctest
julia> X=Polymake.fulton.NormalToricVariety(INPUT_RAYS=[1 0; 1 2],INPUT_CONES=[[0,1]]);

julia> F=addStackStructure(X,[2,3]);

julia> findStackyBarycenter([1,2],F)
[ 5 ,  6 ]
```
"""
function findStackyBarycenter(s::Union{AbstractSet,AbstractVector},SX::StackyFan)
    rayMatrix=convert(Array{Int64,2}, Array(Polymake.common.primitive(SX.fan.RAYS)))
    # Multiply the rays by their stacky values
    stackMatrix = diagm(stackyWeights(SX)) * rayMatrix
    rays = rowMinors(stackMatrix, s)
    dim=size(rays,2)
    bary=zeros(Int64,dim,1)
    for i in 1:size(rays,1)
        bary+=rays[i,:]
    end
    return vec(bary)
end

"""
    findStackyRayMatrix(::StackyFan)

    Outputs the ray matrix of the given stacky fan, such that each primitive ray is multiplied by its stacky weights.

# Examples
```jldoctest
julia> X=Polymake.fulton.NormalToricVariety(INPUT_RAYS=[1 0; 2 3; 1 5],INPUT_CONES=[[0,1],[1,2]]);

julia> F=addStackStructure(X,[1,2,3]);

julia> findStackyRayMatrix(F)
3×2 Matrix{Int64}:
 1   0
 4   6
 3  15
```
"""
function findStackyRayMatrix(sf::StackyFan)
    rayMatrix=convert(Array{Int64,2},Array(Polymake.common.primitive(sf.fan.RAYS)))
    sW=stackyWeights(sf)
    return sW .* rayMatrix
end

"""
    getConeRank(::AbstractMatrix,::AbstractVector)

Takes a matrix and a vector containing indices corresponding to rows of a matrix, and calculates the rank of the matrix consisting only of those rows.

#Examples
```jldoctest makeSmoothWithDependencies
julia> v=[1,2]

julia> M=[0 1; 1 1; 1 0]

julia> getConeRank(v,M)
2
```
"""
function getConeRank(coneRayIndices::AbstractVector, rayMatrix::AbstractMatrix)
    coneRays = rowMinors(rayMatrix,coneRayIndices)
    return rank(Matrix(coneRays))
end

"""
    getDimension(::Polymake.BigObjectAllocated)

Returns the ambient dimension of a normal toric variety.

#Examples
```jldoctest makeSmoothWithDependencies
julia> X=Polymake.fulton.NormalToricVariety(INPUT_RAYS=[1 0 0;1 2 0;0 0 1;0 1 0; 1 1 1],INPUT_CONES=[[0,1,2],[0,2,3,4]])

julia> getDimension(X)
3
```
"""
function getDimension(X)
    return size(X.RAYS, 2)
end

"""
    getConeFaces(::Polymake.BigObjectAllocated,::AbstractVector,::AbstractMatrix)

Takes a fan, its ray matrix, and a vector corresponding to one of its cones, and returns a list of maximal strict faces of that cone.

#Examples
```jldoctest makeSmoothWithDependencies
julia> X=Polymake.fulton.NormalToricVariety(INPUT_RAYS=[1 0 0; 1 1 0; 1 0 1; 1 1 1],INPUT_CONES=[[0,1,2,3]])

julia> getConeFaces(X,[1,2,3,4],Array(X.RAYS))
[[ 1 ,  2 ], [ 1 ,  3 ], [ 3 ,  4 ], [ 2 ,  4 ]]
```
"""
function getConeFaces(fan::Polymake.BigObjectAllocated,cone::AbstractVector,rayMatrix::AbstractMatrix)
    lattice = fan.HASSE_DIAGRAM
    faces = @Polymake.convert_to Array{Set{Int}} lattice.FACES
    cone_faces=[]
    c = rank(Array(rowMinors(rayMatrix, cone))) - 1
    rank_c_subcone_indices = @Polymake.convert_to Array{Int} Polymake.graph.nodes_of_rank(lattice,c)
    rank_c_subcones = [faces[i + 1] for i in rank_c_subcone_indices]
    for subcone in rank_c_subcones
        new_cone = [i+1 for i in subcone]
        if all((i -> i in cone).(new_cone))
            push!(cone_faces, new_cone)
        end
    end
    return cone_faces
end

"""
    toric_blowup(::Union{AbstractSet,AbstractVector},::Polymake.BigObjectAllocated,::AbstractVector)

    Takes a normal toric variety X, a set s corresponding to a subset of rays of X, and a (optional) polymake vector, v, blow up X at v. If v is not provided, blow up X at the barycenter of s.

# Examples
```jldoctest
julia> X=Polymake.fulton.NormalToricVariety(INPUT_RAYS=[1 0; 1 2],INPUT_CONES=[[0,1]]);

julia> B=toric_blowup([0,1],X,[1,1]);

julia> B.INPUT_RAYS
pm::Matrix<pm::Rational>
1 0
1 2
1 1

julia> B.INPUT_CONES
pm::IncidenceMatrix<pm::NonSymmetric>
{0 2}
{1 2}
```
"""
function toric_blowup(s, X, v)
    s = [i + 1 for i in s]
    if v==nothing
        v=findBarycenter(s,X)
    end
    if size(v,2)==1
         v=transpose(v)
    end
    coneList = convertIncidenceMatrix(X.MAXIMAL_CONES)
    # Extracting the indices of all the cones in X that contains the set of rays s
    starIndex = findall((t) -> all(((i) -> i in t).(s)), coneList)
    star = [coneList[i] for i in starIndex]
    rayMatrix = X.RAYS

    lattice = X.HASSE_DIAGRAM
    faces = @Polymake.convert_to Array{Set{Int}} lattice.FACES

    # Get all the subcones of X that is contained in one of the cones in star and has the same rank
    clStar = []
    # Iterate over star
    for t in star
        # Get the rank of t
        c = rank(Array(rowMinors(rayMatrix, t))) - 1
        # Get all the subcones of X with rank c
        rank_c_subcone_indices = @Polymake.convert_to Array{Int} Polymake.graph.nodes_of_rank(lattice,c)
        rank_c_subcones = [faces[i + 1] for i in rank_c_subcone_indices]
        # Iterate over rank_c_subcones, and put the cones that is contained in t into clStar
        for cone in rank_c_subcones
            new_cone = [i+1 for i in cone]
            if all((i -> i in t).(new_cone))
                push!(clStar, new_cone)
            end
        end
    end
    clStar = unique(clStar)

    n = size(rayMatrix, 1) + 1
    # Filter out the cones in star from conelist
    coneList = filter(x -> !(x in star), coneList)

    if length(s) == 1
        # If s consists of a single ray, find all the cones in clStar that does not contain s
        newCones = []
        for t in clStar
            if !(s[1] in t)
                push!(newCones, sort(push!(t, s[1])))
            end
        end
        # return newCones plus coneList
        finalCones = [[i - 1 for i in cone] for cone in append!(coneList, newCones)]
        return Polymake.fulton.NormalToricVariety(INPUT_RAYS = Array(X.RAYS), INPUT_CONES = finalCones)
    end
    newCones = []
    for t in clStar
        # Find all the cones in clStar that does not contain at least one ray in s
        # QUESTION: Why seperate this from the one element case? Any won't work with one element list?
        if any(((i) -> !(i in t)).(s))
            push!(newCones, push!(t, n))
        end
    end
    # return newCones plus coneList
    finalRays = vcat((X.RAYS),v)
    finalCones = [[i - 1 for i in cone] for cone in append!(coneList, newCones)]
    return Polymake.fulton.NormalToricVariety(INPUT_RAYS = finalRays, INPUT_CONES = finalCones)
end

"""
    makeSimplicial(::Polymake.BigObjectAllocated)

Takes in a normal toric variety and returns a simplicial toric variety by subdividing (blowing up) the non-simplicial maximal cones.

#Examples
```jldoctest makeSmoothWithDependencies
julia> X=Polymake.fulton.NormalToricVariety(INPUT_RAYS=[1 0 0;1 1 0;1 0 1;1 1 1],INPUT_CONES=[[0,1,2,3]])

julia> X.SIMPLICIAL
false

julia> makeSimplicial(X).SIMPLICIAL
true

julia> X=Polymake.fulton.NormalToricVariety(INPUT_RAYS=[1 0 0 0;0 1 0 0;0 0 1 0;1 -1 1 0; 1 0 -2 0],INPUT_CONES=[[0,1,2,3],[0,4]])

julia> X.SIMPLICIAL
false

julia> makeSimplicial(X).SIMPLICIAL
true
```
"""
function makeSimplicial(X::Polymake.BigObjectAllocated)
    Y = copy(X)
    while (true)
        # If the initial toric variety is simplicial, the program terminates and returns it.
        if Y.SIMPLICIAL==true
            break
        end
        #Maximal cones and ray matrix
        coneList = convertIncidenceMatrix(Y.MAXIMAL_CONES)
        rayMatrix = Y.RAYS
        badCone = nothing
        for i in 1:size(coneList,1)
            cone = coneList[i]
            if (getConeRank(cone, rayMatrix) != size(cone)[1])
                badCone = cone
            end
        end
        if (badCone == nothing)
            # All cones are linearly independent
            break
        else
            # Find the first ray that is contained in more than one orbit
            # and subdivide at that ray, using toricBlowup

            # Get faces (need to replace this)
            edges = getConeFaces(Y,badCone,rayMatrix)
            # Find the first ray that is contained in more than one orbit
            i = 1
            while count(r->(badCone[i] in r), edges) == 1
                i += 1
            end
            # Subdivide at the cone containing just that ray
            badCone=[i-1 for i in badCone]
            Y = toric_blowup(badCone, Y,nothing)
            #Y = toric_blowup([badCone[i]], Y,nothing)
        end
        # Repeat this process until there are no more bad cones
    end
    return Y
end

"""
    makeSmooth(::Polymake.BigObjectAllocated)

Takes in a normal toric variety X and output a new smooth toric variety by iteratively blowing up. In the language of fans, these blowups are achieved by subdividing non-smooth cones.

#Examples
```jldoctest makeSmoothWithDependencies
julia> X=Polymake.fulton.NormalToricVariety(INPUT_RAYS=[4 -1; 0 1],INPUT_CONES=[[0, 1]])

julia> X.SMOOTH_FAN
false

julia> makeSmooth(X).SMOOTH_FAN
true

julia> makeSmooth(X).INPUT_RAYS
pm::Matrix<pm::Rational>
1 -1/4
0 1
1 0

julia>X=Polymake.fulton.NormalToricVariety(INPUT_RAYS=[1 0 0;1 1 0;1 0 1;1 1 1],INPUT_CONES=[[0,1,2,3]])

julia>X.SMOOTH_FAN
false

julia>makeSmooth(X).SMOOTH_FAN
true

julia>makeSmooth(X).INPUT_RAYS
pm::Matrix<pm::Rational>
1 0 0
1 1 0
1 0 1
1 1 1
2 1 1

julia> X=Polymake.fulton.NormalToricVariety(INPUT_RAYS=[1 0 0;0 1 0;0 0 1;0 -1 -1; -1 0 -1; -2 -1 0],INPUT_CONES=[[0,1,2],[0,1,3],[1,3,4],[1,2,4],[2,4,5],[0,2,5],[0,3,5],[3,4,5]])

julia> X.SMOOTH_FAN
false

julia> makeSmooth(X).SMOOTH_FAN
true
```
"""
function makeSmooth(X::Polymake.BigObjectAllocated)
    Y  = copy(X)
    while(true)
        coneList = convertIncidenceMatrix(Y.MAXIMAL_CONES)
        rayMatrix = Array(Polymake.common.primitive(Y.RAYS))
        k = 1
        # Iterate through the coneList, getting the index of the first cone not smooth
        for coneSet in coneList
            # Getting the number of rays in coneSet
            S=size(coneSet)[1]
            coneRays=rowMinors(rayMatrix,coneSet)
            # Checking whether this cone is smooth
            smoothCheck=Polymake.fan.check_fan_objects(Polymake.polytope.Cone(RAYS=coneRays)).SMOOTH_FAN
            if !smoothCheck
                # If the cone is not simplicial or not smooth, we have found the cone that we need to make smooth
                break
            else
                k+=1
            end
        end
        # At this point, all the cones are smooth. The program terminates.
        if k == size(coneList,1)+1
            break
        end

        # Get the cone that we found to be not smooth
        sigma=coneList[k]
        sigmaRays=slicematrix(rowMinors(rayMatrix,sigma))
        tau=0; tauRays=0; tauCone=0
        # Iterate over the subcones of sigma, finding tau, the smallest one that is not smooth
        for subset in collect(powerset(sigma))
            if size(subset,1) > 1
                S=size(subset)[1]
                subsetRays=rowMinors(rayMatrix,subset)
                subsetCone=Polymake.polytope.Cone(RAYS=subsetRays)
                smoothCheck=Polymake.fan.check_fan_objects(subsetCone).SMOOTH_FAN
                if !smoothCheck
                    tau=subset
                    tauRays=subsetRays
                    tauCone=subsetCone
                    break
                end
            end
        end

        # Getting the Hilbert Basis of tau
        H=slicematrix(Matrix(tauCone.HILBERT_BASIS_GENERATORS[1]))
        rayIndex=0
        # Iterate over the Hilbert Basis, finding the first ray that is not the generator of sigma
        for i in 1:size(H,1)
            if !(H[i] in sigmaRays)
                rayIndex=i
                break
            end
        end
        if rayIndex==0
            # Every Hilbert Basis of tau is a generator of sigma. Make Y simplicial is sufficient to make sigma smooth
            Y=makeSimplicial(Y)
        else
            # blowupRay is not a generator of sigma, blow up tau at blowupRay
            blowupRay=H[rayIndex]
            tau=[i-1 for i in tau]
            Y=toric_blowup(tau,Y,transpose(blowupRay))
        end
    end
    return Y
end

"""
    stackyBlowup(::StackyFan,::Array{Int64,1},::Array{Int64,1})

    Takes a stacky fan sf, a ray excep, and a cone, and subdivides the stacky fan at the given ray. Crucially, the given cone should be the minimal cone containing the exceptional ray. The cone input should be zero-indexed.

#examples
```jldoctest
julia> X=Polymake.fulton.NormalToricVariety(INPUT_RAYS=[1 0; 1 2],INPUT_CONES=[[0,1]]);

julia> F=addStackStructure(X,[2,3]);

julia> stackyWeights(stackyBlowup(F,[0,1],[1,1]))
[ 2 ,  1 ,  3 ]
```
"""
function stackyBlowup(sf::StackyFan, cone::Array{Int64,1}, excep::Array{Int64,1})
    # Express the exceptional ray as a scalar multiple of a primitive ray
    # Use this scalar as the stacky weight in the resulting stacky fan
    G=gcd(excep)
    excep=Polymake.common.primitive(excep)

    # Perform toric blowup at the given ray
    blowup = toric_blowup(cone, sf.fan, excep)
    sf.stacks[encode(excep)] = G

    return(StackyFan(blowup, sf.stacks))
end

"""
    getConesPolymake(sf::StackyFan)

    Returns a list of cones of a stacky fan, with the cones represented as polymake objects.

#Examples
```jldoctest
julia> F=makeStackyFan([1 0; 1 1; 1 2; 1 3],[[0,1],[1,2],[2,3]],[1,1,1,1])

julia> getConesPolymake(F)[1].RAYS
pm::Matrix<pm::Rational>
1 0
1 1
```
"""
function getConesPolymake(sf::StackyFan)
    formatted = convertIncidenceMatrix(sf.fan.MAXIMAL_CONES)
    cones = map((x) -> Polymake.polytope.Cone(
        INPUT_RAYS=sf.fan.RAYS[x,:]), formatted)
    return(cones)
end

"""
    slicematrix(::AbstractMatrix{<:Number})

    Take a two-dimensional matrix and output a list of its row vectors.

# Examples
```jldoctest
julia> A=[1 2; 3 4];

julia> slicematrix(A)
[[ 1 ,  2 ], [ 3 ,  4 ]]
```
"""
function slicematrix(A::AbstractMatrix{<:Number})
    return [A[i, :] for i in 1:size(A,1)]
end

"""
    rowMinors(::AbstractMatrix{<:Number},::Union{AbstractSet,AbstractVector})

    Identical to slicematrix, except only returns row vectors indexed by a set S.

# Examples
```jldoctest
julia> A=[1 2 3;4 5 6; 7 8 9];

julia> S=Set([1,3]);

julia> rowMinors(A,S)
2×3 LinearAlgebra.Transpose{Int64,Array{Int64,2}}:
 1  2  3
 7  8  9
```
"""
function rowMinors(A::AbstractMatrix{<:Number},S::Union{AbstractSet,AbstractVector})
    outList=[]
    slices=slicematrix(A)
    for i in 1:size(slices,1)
        if i in S
            append!(outList,[slices[i]])
        end
    end
    return Array(transpose(hcat(outList...)))
end

"""
    convertIncidenceMatrix(::Polymake.IncidenceMatrixAllocated{Polymake.NonSymmetric})

    Takes a Polymake incidence matrix (e.g., the output of X.MAXIMAL_CONES for a toric variety X) and outputs a list of vectors, with each vector recording the indices marked on a given row of the incidence matrix.

# Examples
```jldoctest
julia> X=Polymake.fulton.NormalToricVariety(INPUT_RAYS=[1 0;1 1; 0 1],INPUT_CONES=[[0,1],[1,2]]);

julia> M=X.MAXIMAL_CONES;

julia> convertIncidenceMatrix(M)
[[ 1 ,  2 ], [ 2 ,  3 ]]
```
"""
function convertIncidenceMatrix(A::Polymake.IncidenceMatrixAllocated{Polymake.NonSymmetric})
    A=Array(A)
    dim1=size(A,1)
    dim2=size(A,2)
    out=[]
    for i in 1:dim1
        members=[]
        for j in 1:dim2
            if A[i,j]==true
                append!(members,j)
            end
        end
        append!(out,[members])
    end
    return convert(Array{Array{Int64,1},1}, out)
end

"""
    coneMultiplicity(C::Polymake.BigObjectAllocated)

    Returns the multiplicity of a polyhedral cone (inputted as a Polymake object): here, the multiplicity is defined as the index of the sublattice generated by the rays of the cone, inside the full integer lattice contained in the linear subspace generated by the edges of the cone.

# Examples
```jldoctest

julia> C=Polymake.polytope.Cone(INPUT_RAYS=[1 0; 1 2]);

julia> coneMultiplicity(C)
2
```
"""
function coneMultiplicity(C::Polymake.BigObjectAllocated)
    A=Polymake.common.primitive(C.RAYS)
    M=matrix(ZZ,[fmpz.(y) for y in A])
    SNF=Nemo.snf(M)
    mult=1
    for i in 1:size(SNF,1)
        mult*=SNF[i,i]
    end
    return mult
end

"""
    getMultiplicity(::Array{Int64,1},::Array{Int64,2})

    Same functionality as coneMultiplicity, but calculates the cone rays as a subset of the columns of a ray matrix rather than from a Polymake cone object.

# Examples
```jldoctest
julia> getMultiplicity([1,2],[1 0; 1 2; 1 3])
2
```
"""
function getMultiplicity(cone::Array{Int64,1},rayMatrix::Array{Int64,2})
    A=rowMinors(rayMatrix,cone)
    M=matrix(ZZ,[fmpz.(y) for y in A])
    SNF=Nemo.snf(M)
    mult=1
    for i in 1:size(SNF,1)
        mult*=SNF[i,i]
    end
    return mult
end

"""
    coneConvert(::abstractVector{Int64},::abstractMatrix{Int64})

    Takes a matrix where the columns represent rays, and a list of indices, and forms a Polymake cone object generated by the rays corresponding to those indices.

# Examples
```jldoctest

julia> typeof(coneConvert([1, 2, 4],[1 0 0; 0 1 0; 0 0 1; 1 1 1]))
Polymake.BigObjectAllocated
```
"""
function coneConvert(cone::Array{Int64,1},rayMatrix::Array{Int64,2})
    coneRays=rowMinors(rayMatrix,cone)
    C=Polymake.polytope.Cone(RAYS=coneRays)
    return C
end

"""
    getCones(X::Polymake.BigObjectAllocated)

    Returns all the cones of a fan X as a list of lists, with each interior list containing the indices of the rays generating a given cone.

# Examples
```jldoctest
julia> X=Polymake.fulton.NormalToricVariety(INPUT_RAYS=[1 0 0; 1 1 0; 1 0 1; 1 1 ],INPUT_CONES=[[0,1,2,3]]);

julia> getCones(X)
[[ 0 ,  1 ,  2 ,  3 ], [ 0 ,  1 ], [ 0 ,  2 ], [ 2 ,  3 ], [ 1 ,  3 ], [ 0 ], [ 1 ], [ 2 ], [ 3 ]]
```
"""
function getCones(X::Polymake.BigObjectAllocated)
    lattice=X.HASSE_DIAGRAM
    faces=@Polymake.convert_to Array{Set{Int}} lattice.FACES
    out=Array{Int64,1}[]
    for i in 2:(size(faces,1)-1)
        newface=Array(@Polymake.convert_to Array{Int} faces[i])
        push!(out,[i+1 for i in newface])
    end
    return out
end

"""
    findFaceContainingRay(::Polymake.BigObjectAllocated,::Array{Int64, 1})

Given a cone and a ray, finds a maximal proper face of the cone containing that ray.

# Examples
```jldoctest
julia> C=Polymake.polytope.Cone(INPUT_RAYS=[1 0 0; 1 0 1; 1 1 0; 1 1 1]);

julia> findFaceContainingRay(C,[1,1,1]).RAYS
pm::Matrix<pm::Rational>
1 0 1
1 1 1
```
"""
function findFaceContainingRay(C::Polymake.BigObjectAllocated,v::Array{Int64, 1})
    rayMatrix=convert(Array{Int64,2},Array(Polymake.common.primitive(C.RAYS)))
    faces=convertIncidenceMatrix(C.RAYS_IN_FACETS)
    if faces==[[]]
        return nothing
    end
    for face in faces
        faceCone=coneConvert(face,rayMatrix)
        if Polymake.polytope.contains(faceCone, v)
            return faceCone
        end
    end
    return nothing
end

"""
    findMinimalCone(::Polymake.BigObjectAllocated,::Array{Int64,1})

Given a fan and a ray, finds the minimal cone fo the fan containing the ray.

# Examples
```jldoctest
julia> X=Polymake.fulton.NormalToricVariety(INPUT_RAYS=[1 0 0 0; 1 0 1 0; 1 1 0 0; 1 1 1 0; 0 0 0 1],INPUT_CONES=[[0,1,2,3,4]]);

julia> findMinimalCone(X,[1,2,0,0])
[1, 3]
```
"""
function findMinimalCone(X::Polymake.BigObjectAllocated,v::Array{Int64, 1})
    rayMatrix=convert(Array{Int64,2},Array(Polymake.common.primitive(X.RAYS)))
    cones=convertIncidenceMatrix(X.MAXIMAL_CONES)
    #print(cones)
    startCone=nothing
    for cone in cones
        #print(cone)
        polyCone=coneConvert(cone,rayMatrix)
        if Polymake.polytope.contains(polyCone, v)
            startCone=polyCone
        end
    end
    if startCone==nothing
        error("The given ray is not contained in any cone of the fan.")
    end
    currentCone=startCone
    while(true)
        nextCone=findFaceContainingRay(currentCone,v)
        if nextCone==nothing
            break
        end
        currentCone=nextCone
    end
    currentRays=slicematrix(Array(currentCone.RAYS))
    fanRays=slicematrix(rayMatrix)
    indices=findall(x->x in currentRays,fanRays)
    return indices
end

"""
    starSubdivision(::Polymake.BigObjectAllocated,::Array{Int64,1})

    Blows up the input fan at the input ray. A wrapper integrating findMinimalCone and toric_blowup.

# Examples
```jldoctest
julia> X=Polymake.fulton.NormalToricVariety(INPUT_RAYS=[1 0 0 0; 1 0 1 0; 1 1 0 0; 1 1 1 0; 0 0 0 1],INPUT_CONES=[[0,1,2,3,4]]);

julia> starSubdivision(X,[2,1,0,0]).RAYS
pm::Matrix<pm::Rational>
1 0 0 0
1 0 1 0
0 0 0 1
1 1/2 0 0
1 1 0 0
1 1 1 0
```
"""
function starSubdivision(X::Polymake.BigObjectAllocated, v::Array{Int64, 1})
    minimalCone = findMinimalCone(X, v)
    s = [i - 1 for i in minimalCone]
    v = transpose(v)
    return toric_blowup(s, X, v)
end

"""
    convertToIncidence(v::Array{Int64,1},l::Int64)

Returns a vector of length l, with entries of 1 indexed by v and entries of 0 everywhere else.

# Examples
```jldoctest
julia> convertToIncidence([2,3,5],6)
[ 0 , 1 , 1 , 0 , 1 , 0 ]
```
"""
function convertToIncidence(v::Array{Int64,1},l::Int64)
    out=[]
    for j in 1:l
        if j in v
            append!(out,1)
        else
            append!(out,0)
        end
    end
    return out
end

"""
    interiorPoints(::Polymake.BigObjectAllocated)

    Finds all interior lattice points contained in the fundamental region of a given cone. When multiple interior lattice points lie along the same ray, only the point closest to the origin is returned. Notably,

# Examples
```jldoctest
julia> C=Polymake.polytope.Cone(INPUT_RAYS=[1 2; 2 1]);

julia> interiorPoints(C)
[[ 1 ,  1 ]]
```
"""
function interiorPoints(C::Polymake.BigObjectAllocated)
    rayMatrix=Array(Polymake.common.primitive(C.RAYS))
    l=size(rayMatrix,1)
    dim=size(rayMatrix,2)
    if rank(rayMatrix)<l
        print(rayMatrix)
        error("Input cone is not simplicial.")
    end
    subsets=collect(powerset([1:l;]))
    vertices=[]
    for elt in subsets #vertices of the fundamental region are in correspondence with subsets of the generators of the cone, by summing the generators in a subset to obtain a vertex
        vert=zeros(Polymake.Rational,1,dim)
        for i in 1:l
            if i in elt
                vert+=rayMatrix[[i],:]
            end
        end
        append!(vertices,[vert])
    end
    V=vcat(vertices...)
    VH=hcat(ones(Polymake.Rational,size(V,1)),V)
    P=Polymake.polytope.Polytope(POINTS=VH) #make a Polymake polytope object from the vertices of the fundamental region found in the last step
    if size(P.INTERIOR_LATTICE_POINTS,1)==0
        return nothing
    end
    intPoints=Array(P.INTERIOR_LATTICE_POINTS)[:,2:(dim+1)] #find all the interior lattice points
    validPoints=[]
    #return intPoints
    for i in 1:size(intPoints,1) #throw out all points that are integer multiples of other points
        point=intPoints[i,:]
        if gcd(point)==1
            append!(validPoints,[point])
        end
    end
    return validPoints
end

"""
    coneRayDecomposition(::Array{Int64,1},::Array{Int64,2},::Array{Int64,1},::Array{Int64,1})

    This function takes in a cone (a vector of indices of cone generators in rayMatrix), a ray, and a stacky structure for rayMatrix. It first multiplies all generators of the cone by their stacky values, and then finds an expression for the ray as a sum of these stacky generators. The output is a vector of coefficients of the above representation in terms of the rays in rayMatrix, with zeros as coefficients for all rays not in the given cone.

# Examples
```jldoctest
julia> coneRayDecomposition([1,2,3],[3 5 7; 8 16 9;2 1 3;1 1 1],[2,2,3],[1,1,1,1])
[ 6 ,  5 ,  52 ,  0 ]
```
"""
function coneRayDecomposition(cone,rayMatrix,ray,stack)
    stackMatrix=diagm(stack)*rayMatrix # multiply all rays by stack values
    coneRays=rowMinors(stackMatrix,cone) # find the (stackified) rays in the given cone
    if rank(coneRays)<size(coneRays,1)
        error("The given cone is not simplicial.")
    end
    B=Polymake.common.null_space(hcat(transpose(coneRays),-ray)) # Express the input ray in terms of the stackified cone generators
    N=convert(Array{Int64,1},vec(B))
    if size(N,1)==0
        error("The given ray is not in the span of the cone generators.")
    end
    if N[end]<0 #since the nullspace has arbitrary sign, fix it so the coefficients are all positive
        N*=-1
    end
    pop!(N)
    out=zeros(Int64,size(rayMatrix,1))
    for i in 1:size(N,1)#rewrite the coefficients vector in terms of all the rays in rayMatrix, by padding with zeros when appropriate.
        out[cone[i]]=N[i]
    end
    return out
end

"""
=========VISUALIZATION FUNCTIONALITY==========
"""

using Plots

"""
    plot3dSimpCone(::Array{Array{Int64,1},1})

    Give a list of three 3-dimensional vectors, plots the polygon defined by those vectors and the origin via Plots.jl.

# Examples
```jldoctest
julia> plot3dSimpCone([[1,0,0],[0,1,0],[0,0,1]]);
```
"""
function plot3dSimpCone(L::Array{Array{Int64,1},1})
    l=size(L,1)
    if l==3
        (x1,y1,z1)=L[1]
        (x2,y2,z2)=L[2]
        (x3,y3,z3)=L[3]
        X=[0,x1,0,x2,0,x3,x2,x1,x3]
        Y=[0,y1,0,y2,0,y3,y2,y1,y3]
        Z=[0,z1,0,z2,0,z3,z2,z1,z3]
        plot!(X,Y,Z)
    elseif l==2
        (x1,y1,z1)=L[1]
        (x2,y2,z2)=L[2]
        X=[0,x1,x2,0]
        Y=[0,y1,y2,0]
        Z=[0,z1,z2,0]
        plot!(X,Y,Z)
    elseif l==1
        (x1,y1,z1)=L[1]
        X=[0,x1]
        Y=[0,y1]
        Z=[0,z1]
        plot!(X,Y,Z)
    else
        error("Input cone is empty or not simplicial.")
    end
end

"""
    plot2dCone(::Array{Array{Int64,1},1})

    Given a list of two 2-dimesnional vectors, plots the polygon defined by those vectors and the origin via Plots.jl.

# Examples
```jldoctest
julia> plot2dCone([[1,0],[0,1]]);
```
"""
function plot2dCone(L::Array{Array{Int64,1},1})
    l=size(L,1)
    if l==2
        (x1,y1)=L[1]
        (x2,y2)=L[2]
        S=Shape([(0,0),(x1,y1),(x2,y2)])
        plot!(S)
    elseif l==1
        (x1,y1)=L[1]
        X=[0,x1]
        Y=[0,y1]
        plot!(X,Y)
    else
        error("Input cone is empty or not simplicial.")
    end
end

"""
    showSimpFan(X::Polymake.BigObjectAllocated)

    Plots the input simplicial fan as a collection of polygons or polyhedra defined by the input rays and the origin via Plots.jl.

# Examples
```jldoctest
julia> X=Polymake.fulton.NormalToricVariety(INPUT_RAYS=[1 0; 2 5],INPUT_CONES=[[0,1]]);

julia> showSimpFan(X);
```
"""
function showSimpFan(X::Polymake.BigObjectAllocated)
    plot()
    rayMatrix=convert(Array{Int64,2},Array(Polymake.common.primitive(X.RAYS)))
    dim=size(rayMatrix,2)
    if !(dim in [2,3])
         println("Fan visualization is only supported in dimensions 2 and 3.")
    end
    maxCones=convertIncidenceMatrix(X.MAXIMAL_CONES)
    for cone in maxCones
        coneRays=slicematrix(rowMinors(rayMatrix,cone))
        if dim==2
            plot2dCone(coneRays)
        elseif dim==3
            plot3dSimpCone(coneRays)
        end
    end
    plot!(legend=false)
end

"""
    showSimpStackyFan(::StackyFan;::Bool=true)

    Extends the functionality of showSimpFan to stacky simplicial fans. If the stackypoints input is set to true, a red dot is shown at the location of each primitive vector multiplied by its stacky weight.

# Examples
```jldoctest
julia> X=Polymake.fulton.NormalToricVariety(INPUT_RAYS=[1 0; 2 5],INPUT_CONES=[[0,1]]);

julia> F=addStackStructure(X,[1,2]);

julia> showSimpStackyFan(F);
```
"""
function showSimpStackyFan(F::StackyFan;stackypoints::Bool=true)
    X=F.fan
    plot()
    rayMatrix=convert(Array{Int64,2},Array(Polymake.common.primitive(X.RAYS)))
    dim=size(rayMatrix,2)
    if !(dim in [2,3])
         println("Fan visualization is only supported in dimensions 2 and 3.")
    end
    maxCones=convertIncidenceMatrix(X.MAXIMAL_CONES)
    for cone in maxCones
        coneRays=slicematrix(rowMinors(rayMatrix,cone))
        if dim==2
            plot2dCone(coneRays)
        elseif dim==3
            plot3dSimpCone(coneRays)
        end
    end
    if stackypoints==true
        stackyPoints=slicematrix(findStackyRayMatrix(F))
        for point in stackyPoints
            if dim==2
                (x,y)=point
                plot!([x],[y],shape=:circle,color=:red)
            elseif dim==3
                (x,y,z)=point
                plot!([x],[y],[z],shape=:circle,color=:red)
            end
        end
    end
    plot!(legend=false)
end

"""
    coneVectorOrder(::Polymake.BigObjectAllocated)

    Takes a polyhedral cone in 2 or 3 dimensions and outputs a list of its defining rays, arranged in counterclockwise order around the exterior of the cone.

# Examples
```jldoctest
julia> C=Polymake.polytope.Cone(INPUT_RAYS=[1 0 0; 1 0 1; 1 1 0; 1 1 1]);

julia> coneVectorOrder(C)
[[1
, 1, 0], [1, 1, 1], [1, 0, 1], [1, 0, 0]]
```
"""
function coneVectorOrder(C::Polymake.BigObjectAllocated)
    rayList=slicematrix(convert(Array{Int64,2},Array(Polymake.common.primitive(C.RAYS))))
    orderedRayList=Array{Int64,1}[]
    ordering=Array(C.RIF_CYCLIC_NORMAL[1])
    ordering=[i+1 for i in ordering]
    for i in ordering
        push!(orderedRayList,rayList[i])
    end
    return orderedRayList
end

"""
    plot3dCone(::Array{Array{Int64,1},1})

    Plots the 3-dimensional cone defined by the given list of rays. The input rays are assumed to be in counterclockwise or clockwise order around the cone; coneVectorOrder() can be used to obtain this ordering from an arbitrary cone.

# Examples
```jldoctest
julia> plot3dCone([[1,1,0],[1,1,1],[1,0,1],[1,0,0]]);
```
"""
function plot3dCone(L::Array{Array{Int64,1},1})
    l=size(L,1)
    X=Int64[0]
    Y=Int64[0]
    Z=Int64[0]
    for i in 1:l
        push!(X,L[i][1])
        push!(Y,L[i][2])
        push!(Z,L[i][3])
        push!(X,0)
        push!(Y,0)
        push!(Z,0)
    end
    for i in 1:l
        push!(X,L[i][1])
        push!(Y,L[i][2])
        push!(Z,L[i][3])
    end
    push!(X,L[1][1])
    push!(Y,L[1][2])
    push!(Z,L[1][3])
    plot!(X,Y,Z)
end

"""
    showFan(X::Polymake.BigObjectAllocated)

    Plots the input fan as a collection of polygons or polyhedra defined by the input rays and the origin via Plots.jl.

# Examples
```jldoctest
julia> X=Polymake.fulton.NormalToricVariety(INPUT_RAYS=[1 0 0; 1 0 1; 1 1 0; 1 1 1],INPUT_CONES=[[0,1,2,3]]);

julia> showFan(X);
```
"""
function showFan(X::Polymake.BigObjectAllocated)
    plot()
    rayMatrix=convert(Array{Int64,2},Array(Polymake.common.primitive(X.RAYS)))
    dim=size(rayMatrix,2)
    if !(dim in [2,3])
         println("Fan visualization is only supported in dimensions 2 and 3.")
    end
    maxCones=convertIncidenceMatrix(X.MAXIMAL_CONES)
    for cone in maxCones
        if dim==2
            coneRays=slicematrix(rowMinors(rayMatrix,cone))
            plot2dCone(coneRays)
        elseif dim==3
            C=coneConvert(cone,rayMatrix)
            plot3dCone(coneVectorOrder(C))
        end
    end
    plot!(legend=false)
end

"""
    showStackyFan(::StackyFan;::Bool=true)

    Extends the functionality of showFan to stacky fans. If the stackypoints input is set to true, a red dot is shown at the location of each primitive vector multiplied by its stacky weight.

# Examples
```jldoctest
julia> X=Polymake.fulton.NormalToricVariety(INPUT_RAYS=[1 0 0; 1 0 1; 1 1 0; 1 1 1],INPUT_CONES=[[0,1,2,3]]);

julia> F=addStackStructure(X,[1,2,2,5]);

julia> showStackyFan(F);
```
"""
function showStackyFan(F::StackyFan;stackypoints::Bool=true)
    X=F.fan
    plot()
    rayMatrix=convert(Array{Int64,2},Array(Polymake.common.primitive(X.RAYS)))
    dim=size(rayMatrix,2)
    if !(dim in [2,3])
         println("Fan visualization is only supported in dimensions 2 and 3.")
    end
    maxCones=convertIncidenceMatrix(X.MAXIMAL_CONES)
    for cone in maxCones
        if dim==2
            coneRays=slicematrix(rowMinors(rayMatrix,cone))
            plot2dCone(coneRays)
        elseif dim==3
            C=coneConvert(cone,rayMatrix)
            plot3dCone(coneVectorOrder(C))
        end
    end
    if stackypoints==true
        stackyPoints=slicematrix(findStackyRayMatrix(F))
        for point in stackyPoints
            if dim==2
                (x,y)=point
                plot!([x],[y],shape=:circle,color=:red)
            elseif dim==3
                (x,y,z)=point
                plot!([x],[y],[z],shape=:circle,color=:red)
            end
        end
    end
    plot!(legend=false)
end


module NumericPredicates

import Base: IEEEFloat,
             iszero, isone,
             isinteger, iseven, isodd,
             isreal, isfinite, isinf, isnan, issubnormal

export negone, two, half,
       isnegone, istwo, ishalf,
       ispos, isneg, isnonneg, isnonpos,
       isfloat, isrational, iscomplex, iscomplexint, iscomplexfloat,
       isfractional, ismixed, isextended, nextint, prevint,

 using IEEEFloats

function nextint(x::T) where {T<:Union{Signed,Unsigned}}
     x === typemax(T) && throw(ErrorException("OverflowError($x)"))
     x += one(T)
     return x
end

function prevint(x::T) where {T<:Union{Signed,Unsigned}}
     x === typemin(T) && throw(ErrorException("UnderflowError($x)"))
     x -= one(T)
     return x
end

function nextint(x::T) where {T<:IEEEFloat}
     x = floor(x)
     x === floatintmax(T) && throw(ErrorException("OverflowError($x)"))
     x += one(T)
     return x
end

function prevint(x::T) where {T<:IEEEFloat}
     x = ceil(x)
     x === floatintmin(T) && throw(ErrorException("UnderflowError($x)"))
     x -= one(T)
     return x
end

function nextint(x::T) where {T<:Real}
     x = floor(x)
     z = x + one(T)
     z === x && throw(ErrorException("OverflowError($x)"))
     return z
end

function prevint(x::T) where {T<:Real}
     x = ceil(x)
     z = x - one(T)
     z === x && throw(ErrorException("UnderflowError($x)"))
     return z
end

isfractional(x::T) where {T<:Union{Signed,Unsigned}} = false
isfractional(x::T) where {T<:Real} = -one(T) < x < one(T)
ismixed(x::T) where {T<:Union{Signed,Unsigned}} = false
ismixed(x::T) where {T<:AbstractFloat} = abs(x) > one(T) && !isinteger(x)

end # module NumericPredicates


#Figure 3.2
using Plots, SparseArrays, JLD2, LinearAlgebra, MultiscaleGraphSignalTransforms
include("../../../../src/utils.jl")

JLD2.@load "../data/Toronto.jld2"
tmp1 = toronto["G"]
G=GraphSig(tmp1["W"],xy=tmp1["xy"],f=tmp1["f"],name =tmp1["name"],plotspecs = tmp1["plotspecs"])

G = Adj2InvEuc(G)
GP = partition_tree_fiedler(G,:Lrw)
dmatrix = ghwt_analysis!(G, GP=GP)

N = length(G.f)

j = 1
loc = 1
BS = bs_level(GP, j)
dvec = zeros(N,1)
dvec[loc,1] = 1
(f, GS) = ghwt_synthesis(dvec, GP, BS, G)
GraphSig_Plot(GS, linewidth = 1., markersize = 4., markerstrokealpha =0., notitle = true, clim = (-0.1, 0.1))
plot!(axis = false,colorbar = false)
k,l = BS.levlist[loc][1], BS.levlist[loc][2]
print(j," ",(rs_to_region(GP.rs, GP.tag))[k,l]," ",GP.tag[k,l])
#savefig("Toronto100.pdf")

j = 2
loc = 2
BS = bs_level(GP, j)
dvec = zeros(N,1)
dvec[loc,1] = 1
(f, GS) = ghwt_synthesis(dvec, GP, BS, G)
GraphSig_Plot(GS, linewidth = 1., markersize = 4., markerstrokealpha =0., notitle = true, clim = (-0.1, 0.1))
plot!(axis = false,colorbar = false)
k,l = BS.levlist[loc][1], BS.levlist[loc][2]
print(j," ",(rs_to_region(GP.rs, GP.tag))[k,l]," ",GP.tag[k,l])
#savefig("Toronto201.pdf")

j = 5
loc = 100
BS = bs_level(GP, j)
dvec = zeros(N,1)
dvec[loc,1] = 1
(f, GS) = ghwt_synthesis(dvec, GP, BS, G)
GraphSig_Plot(GS, linewidth = 1., markersize = 4., markerstrokealpha =0., notitle = true, clim = (-0.1, 0.1))
plot!(axis = false,colorbar = false)
k,l = BS.levlist[loc][1], BS.levlist[loc][2]
print(j," ",(rs_to_region(GP.rs, GP.tag))[k,l]," ",GP.tag[k,l])
#savefig("Toronto5134.pdf")


struct Constraint{op,L,R}
    lhs::L
    rhs::R
end

Constraint{op}(lhs::L, rhs::R) where {op, L, R} = Constraint{op,L,R}(lhs, rhs)

function print_op(op)
    (op == ≠) && return "≠"
    (op == ~) && return "="
end

Base.show(io::IO, c::Constraint{op}) where {op} = print(io, c.lhs, " ", print_op(op), " ", c.rhs)

function clauses(c::Constraint{!=})

    # ¬(x[i] ∧ y[i])  -- can't take same value

    [¬(c.lhs[i]) ∨ ¬(c.rhs[i]) for i in 1:min(max(c.lhs), max(c.rhs))]
end

function clauses(c::Constraint{~})
    i = c.rhs  # integer
    [c.lhs[i]]  # that Boolean is positive
end

Base.:≠(x::BoundedInteger, y::BoundedInteger) = Constraint{≠}(x, y)
Base.:~(x::BoundedInteger, y::Int) = Constraint{~}(x, y)

macro integer(name, k)
    quoted = Meta.quot(name)

    :($(esc(name)) = BoundedInteger{$k}($quoted))
end

make_vector(name, T, num) = [T(Symbol(name, subscript(i))) for i in 1:num]


"""
    local_interaction(::AbstractFockAddress, u)
    local_interaction(::AbstractFockAddress, ::AbstractFockAddress, v)

Return the sum of (mode-wise) local interactions ``\\frac{u}{2} \\sum_i n_i(n_i-1)`` of a
single component Fock state, or ``v \\sum_i n_{↑,i} n_{↓,i}`` between two Fock states. For a
multi-component Fock state, return the eigenvalue of

```math
\\frac{1}{2}\\sum_{i, σ, τ} u_{σ,τ} a^†_{σ,i}a^†_{τ,i}a^†_{τ,i}a^†_{σ,i} ,
```

where `u::SMatrix` is a symmetric matrix of interaction constants, `i` is a mode index,
and `σ`, `τ` are component indices.

See also [`BoseFS`](@ref), [`FermiFS`](@ref), [`CompositeFS`](@ref).
"""
local_interaction(b::BoseFS, u) = u * bose_hubbard_interaction(b) / 2
local_interaction(f::FermiFS, _) = 0
local_interaction(f::FermiFS, g::FermiFS, v) = v * count_ones(f.bs & g.bs)
local_interaction(f::FermiFS, b::BoseFS, v) = local_interaction(b, f, v)

function local_interaction(a::BoseFS, b::BoseFS, v)
    occ_a = occupied_modes(a)
    occ_b = occupied_modes(b)

    (n_a, i_a, _), st_a = iterate(occ_a)
    (n_b, i_b, _), st_b = iterate(occ_b)

    acc = 0
    while true
        if i_a > i_b
            # b is behind and needs to do a step
            iter = iterate(occ_b, st_b)
            isnothing(iter) && return acc * v
            (n_b, i_b, _), st_b = iter
        elseif i_a < i_b
            # a is behind and needs to do a step
            iter = iterate(occ_a, st_a)
            isnothing(iter) && return acc * v
            (n_a, i_a, _), st_a = iter
        else
            # a and b are at the same position
            acc += n_a * n_b
            # now both need to do a step
            iter = iterate(occ_a, st_a)
            isnothing(iter) && return acc * v
            (n_a, i_a, _), st_a = iter
            iter = iterate(occ_b, st_b)
            isnothing(iter) && return acc * v
            (n_b, i_b, _), st_b = iter
        end
    end
end
function local_interaction(b::BoseFS, f::FermiFS, v)
    acc = 0
    for (n, i) in occupied_modes(b)
        acc += is_occupied(f, i) * n
    end
    return acc * v
end
function local_interaction(fs::CompositeFS, m)
    _interactions(fs.components, m)
end

"""
    _interaction_col(a, bs::Tuple, us::Tuple)

Sum the local interactions of the Fock state `a` with all states in `bs` using the
interaction constants in `us`. This is used to compute all interactions in the column
below the diagonal of the interaction matrix.
"""
@inline _interaction_col(a, ::Tuple{}, ::Tuple{}) = 0
@inline function _interaction_col(a, (b, bs...), (u, us...))
    return local_interaction(a, b, u) + _interaction_col(a, bs, us)
end

"""
    _interactions(addresses, interaction_matrix)

Compute all pairwise interactions in a tuple of `addresses`. The `interaction_matrix` sets the
intraction strengths.

The code is equivalent to the following.

```julia
acc = 0.0
for (i, a) in enumerate(addresses)
    acc += local_interaction(a, interaction_matrix[i, i])
    for (j, b) in enumerate(addresses[i+1:end])
        acc += local_interaction(a, b, interaction_matrix[i, j])
    end
end
return acc
```

It is implemented recursively to ensure type stability.
"""
@inline _interactions(::Tuple{}, ::SMatrix{0,0}) = 0.0
@inline function _interactions(
    (a, as...)::NTuple{N,AbstractFockAddress}, m::SMatrix{N,N}
) where {N}
    # Split the matrix into the column we need now, and the rest.
    (u, column...) = Tuple(m[:, 1])
    # Type-stable way to subset SMatrix:
    rest = SMatrix{N-1,N-1}(view(m, 2:N, 2:N))

    # Get the self-interaction first.
    self = local_interaction(a, u)
    # Get the interactions for the rest of the row.
    row = _interaction_col(a, as, column)
    # Get the interaction for the rest of the rows.
    return self + row + _interactions(as, rest)
end

###
### HubbardRealSpace
###
"""
    HubbardRealSpace(address; u=ones(C, C), t=ones(C), geometry=PeriodicBoundaries(M,))

Hubbard model in real space. Supports single or multi-component Fock state
addresses (with `C` components) and various (rectangular) lattice geometries
in arbitrary dimensions.

```math
  \\hat{H} = -\\sum_{\\langle i,j\\rangle,σ} t_σ a^†_{iσ} a_{jσ} +
  \\frac{1}{2}\\sum_{i,σ} u_{σ,σ} n_{iσ} (n_{iσ} - 1) +
  \\frac{1}{2}\\sum_{i,σ≠τ}u_{σ,τ} n_{iσ} n_{iτ}
```

## Address types

* [`BoseFS`](@ref): Single-component Bose-Hubbard model.
* [`FermiFS`](@ref): Single-component Fermi-Hubbard model. This address only provides a
  single species of (non-interacting) fermions. You probably want to use [`CompositeFS`](@ref).
* [`CompositeFS`](@ref): For multi-component models.

## Geometries

* [`PeriodicBoundaries`](@ref)
* [`HardwallBoundaries`](@ref)
* [`LadderBoundaries`](@ref)

## Other parameters

* `u`: the on-site interaction parameters. Must be a symmetric matrix. `u[i, j]`
  corresponds to the interaction between the `i`-th and `j`-th component. `u[i, i]`
  corresponds to the interaction of a component with itself. Note that `u[i,i]` must
  be zero for fermionic components.
* `t`: the hopping strengths. Must be a vector of length `C`. The `i`-th element of the
  vector corresponds to the hopping strength of the `i`-th component.

"""
struct HubbardRealSpace{
    C,A,G, # C: components
    # The following need to be type params.
    T<:SVector{C,Float64},
    U<:SMatrix{C,C,Float64},
} <: AbstractHamiltonian{Float64}
    address::A
    u::U # interactions
    t::T # hopping strengths
    geometry::G
end

function HubbardRealSpace(
    address;
    u=ones(num_components(address), num_components(address)),
    t=ones(num_components(address)),
    geometry=PeriodicBoundaries((num_modes(address),))
)
    C = num_components(address)

    # Sanity checks
    if prod(size(geometry)) ≠ num_modes(address)
        error("`geometry` does not have the correct number of sites")
    elseif length(u) ≠ 1 && !issymmetric(u)
        error("`u` must be symmetric")
    elseif length(u) ≠ C * C
        error("`u` must be a $C × $C matrix")
    elseif size(t) ≠ (C,)
        error("`t` must be a vector of length $C")
    elseif address isa BoseFS2C
        error("`BoseFS2C` is not supported for this Hamiltonian, use `CompositeFS`")
    end
    warn_fermi_interaction(address, u)

    u_mat = SMatrix{C,C,Float64}(u)
    t_vec = SVector{C,Float64}(t)
    return HubbardRealSpace{C,typeof(address),typeof(geometry),typeof(t_vec),typeof(u_mat)}(
        address, u_mat, t_vec, geometry,
    )
end

"""
    warn_fermi_interaction(address, u)

Warn if interaction matrix `u` does not make sense for `address`.
"""
function warn_fermi_interaction(address::CompositeFS, u)
    C = num_components(address)
    for c in 1:C
        if address.components[c] isa FermiFS && u ≠ ones(C,C) && u[c,c] ≠ 0
            @warn "component $(c) is fermionic, but was given a self-interaction " *
                "strength of $(u[c,c])" maxlog=1
        end
    end
end
function warn_fermi_interaction(address::FermiFS, u)
    if u ≠ ones(1, 1) && u[1, 1] ≠ 0
        @warn "address is fermionic, but was given a self-interaction " *
            "strength of $(u[1,1])" maxlog=1
    end
end
warn_fermi_interaction(_, _) = nothing

LOStructure(::Type{<:HubbardRealSpace}) = IsHermitian()

function Base.show(io::IO, h::HubbardRealSpace)
    println(io, "HubbardRealSpace(")
    println(io, "  ", starting_address(h), ",")
    println(io, "  u = ", Float64.(h.u), ",")
    println(io, "  t = ", Float64.(h.t), ",")
    println(io, "  geometry = ", h.geometry, ",")
    println(io, ")")
end

starting_address(h::HubbardRealSpace) = h.address
diagonal_element(h::HubbardRealSpace, address) = local_interaction(address, h.u)
diagonal_element(h::HubbardRealSpace{1}, address) = local_interaction(address, h.u[1])

###
### Offdiagonals
###
# This may be an inefficient implementation, but it is not actually used anywhere in the
# main algorithm.
get_offdiagonal(h::HubbardRealSpace, add, i) = offdiagonals(h, add)[i]
num_offdiagonals(h::HubbardRealSpace, add) = length(offdiagonals(h, add))

"""
    HubbardRealSpaceCompOffdiagonals{G,A} <: AbstractOffdiagonals{A,Float64}

Offdiagonals for a single address component. Used with [`HubbardRealSpace`](@ref) model
with a single-component address, or a component of a [`CompositeFS`](@ref).
"""
struct HubbardRealSpaceCompOffdiagonals{G,A} <: AbstractOffdiagonals{A,Float64}
    geometry::G
    address::A
    t::Float64
    length::Int
end

function offdiagonals(h::HubbardRealSpace, comp, add)
    neighbours = num_neighbours(h.geometry)
    return HubbardRealSpaceCompOffdiagonals(
        h.geometry, add, h.t[comp], num_occupied_modes(add) * neighbours
    )
end

Base.size(o::HubbardRealSpaceCompOffdiagonals) = (o.length,)

@inline function Base.getindex(o::HubbardRealSpaceCompOffdiagonals, chosen)
    neighbours = num_neighbours(o.geometry)
    particle, neigh = fldmod1(chosen, neighbours)
    src_index = find_occupied_mode(o.address, particle)
    neigh = neighbour_site(o.geometry, src_index.mode, neigh)

    if neigh == 0
        return o.address, 0.0
    else
        dst_index = find_mode(o.address, neigh)
        new_add, value = excitation(o.address, (dst_index,), (src_index,))
        return new_add, -o.t * value
    end
end

# For simple models with one component.
offdiagonals(h::HubbardRealSpace{1,A}, add::A) where {A} = offdiagonals(h, 1, add)

# Multi-component part
"""
    HubbardRealSpaceOffdiagonals{A,T<:Tuple} <: AbstractOffdiagonals{A,Float64}

Offdiagonals of a [`HubbardRealSpace`](@ref) model with a [`CompositeFS`](@ref) address.
"""
struct HubbardRealSpaceOffdiagonals{A,T<:Tuple} <: AbstractOffdiagonals{A,Float64}
    address::A
    parts::T
    length::Int
end

"""
    get_comp_offdiags(h::HubbardRealSpace, add)

Get offdiagonals of all components of address in a type-stable manner.
"""
@inline function get_comp_offdiags(h::HubbardRealSpace, address)
    return _get_comp_offdiags(address.components, h, Val(1))
end

# All steps of recursive function (should) get inlined, creating a type-stable tuple of
# offdiagonals.
@inline function _get_comp_offdiags((a,as...), h, ::Val{I}) where {I}
    return (offdiagonals(h, I, a), _get_comp_offdiags(as, h, Val(I+1))...)
end
@inline _get_comp_offdiags(::Tuple{}, h, ::Val) = ()

function offdiagonals(h::HubbardRealSpace{C,A}, address::A) where {C,A<:CompositeFS}
    parts = get_comp_offdiags(h, address)
    return HubbardRealSpaceOffdiagonals(address, parts, sum(length, parts))
end

Base.size(o::HubbardRealSpaceOffdiagonals) = (o.length,)

# Becomes type unstable without inline for lots of components. Recursive function is used
# because the type of the result of `o.parts[i]` can not be inferred.
@inline function Base.getindex(o::HubbardRealSpaceOffdiagonals{A}, chosen) where {A}
    return _getindex(o.parts, o.address, chosen, Val(1))
end
@inline function _getindex((p, ps...), address::A, chosen, comp::Val{I}) where {A,I}
    if chosen ≤ length(p)
        new_add, val = p[chosen]
        return BitStringAddresses.update_component(address, new_add, comp), val
    else
        chosen -= length(p)
        return _getindex(ps, address, chosen, Val(I + 1))
    end
end


using LinearAlgebra: Symmetric, isposdef
using Distributions: MvNormal, logpdf

import ..kalman_filter_mv
import ..kalman_smoother_mv

abstract type AbstractLinearGaussian end;

struct LinearGaussian{T<:Any} <: AbstractLinearGaussian
    θ::T
end

function LinearGaussian(; kwargs...)
    LinearGaussian{typeof(kwargs.data)}(kwargs.data)
end

function Base.show(io::IO, ::LinearGaussian)
    print(io, "LinearGaussian()")
end

# Generative Model
function dpr_base(model::AbstractLinearGaussian, Zs)
    logpdf(MvNormal(model.θ.μ0, model.θ.Tau0), Zs)
end

function rpr_base(model::AbstractLinearGaussian, rng)
    rand(
        rng,
        MvNormal(model.θ.states[:, 1], Symmetric(model.θ.variances[:, :, 1])),
        model.θ.K,
    )
end

function dt_base(m::AbstractLinearGaussian, Zs, Zs_new, t)
    μ = m.θ.G(t) * Zs
    Σ = m.θ.Tau(t)
    d = logpdf(MvNormal(Σ), Zs_new - μ)
    @assert !any(isnan.(d))
    return d
end

function dm_base(m::AbstractLinearGaussian, Zs, t)
    yt = view(m.θ.Y, t, :)
    if any(ismissing.(yt))
        return zeros(eltype(Zs), size(Zs, 2))
    else
        d = logpdf(MvNormal(m.θ.Σ(t)), m.θ.F(t) * Zs .- yt)
        @assert !any(isnan.(d))
        return d
    end
end

# Inference Model
function rinit_base(m::AbstractLinearGaussian, rng)
    rand(rng, MvNormal(m.θ.states_sm[:, 1], Symmetric(m.θ.variances_sm[:, :, 1])), m.θ.K)
end

function dinit_base(m::AbstractLinearGaussian, Zs)
    logpdf(MvNormal(m.θ.states_sm[:, 1], Symmetric(m.θ.variances_sm[:, :, 1])), Zs)
end

function calc_prop(m::AbstractLinearGaussian, Zs, t)
    μ =
        m.θ.states_sm[:, t] .+
        (m.θ.Covs_sm[:, :, t-1] / m.θ.variances_sm[:, :, t-1]) *
        (Zs .- m.θ.states_sm[:, t-1])
    Σ = Symmetric(
        m.θ.variances_sm[:, :, t] -
        (m.θ.Covs_sm[:, :, t-1] / m.θ.variances_sm[:, :, t-1]) * m.θ.Covs_sm[:, :, t-1]',
    )

    @assert !any(isnan.(μ))
    @assert !any(isnan.(Σ))

    return μ, Σ
end

function rp_base(m::AbstractLinearGaussian, rng, Zs, t)
    μ, Σ = calc_prop(m, Zs, t)
    @assert isposdef(Σ)
    Zs_new = μ + rand(rng, MvNormal(Σ), size(Zs, 2))
    return Zs_new
end

function dp_base(m::AbstractLinearGaussian, Zs, Zs_new, t)
    μ, Σ = calc_prop(m, Zs, t)
    @assert isposdef(Σ)
    d = logpdf(MvNormal(Σ), Zs_new - μ)
    return d
end

function dpre_base(m::AbstractLinearGaussian, Zs, t)
    d_smooth =
        logpdf(MvNormal(m.θ.states_sm[:, t-1], Symmetric(m.θ.variances_sm[:, :, t-1])), Zs)
    d_filt = logpdf(MvNormal(m.θ.states[:, t-1], Symmetric(m.θ.variances[:, :, t-1])), Zs)

    d = (d_smooth - d_filt) .+ sum(@view(m.θ.lls[t:end]))
    @assert !any(isnan.(d))
    return d
end

function make_kalman_proposal(F0, F, G0, G, Tau, Σ, μ0, Tau0, Y)
    ll, states, variances, lls = kalman_filter_mv(F0, F, G0, G, Σ, Tau, μ0, Tau0, Y)
    states_sm, variances_sm, Covs_sm = kalman_smoother_mv(F0, F, G0, G, Σ, Tau, μ0, Tau0, Y)
    return @NT(ll, states, variances, lls, states_sm, variances_sm, Covs_sm)
end

function lg_smc_model(K, F0, F, G0, G, Tau, Σ, μ0, Tau0, Y)
    ll, states, variances, lls, states_sm, variances_sm, Covs_sm =
        make_kalman_proposal(F0, F, G0, G, Tau, Σ, μ0, Tau0, Y)

    m = LinearGaussian(
        F = F,
        G = G,
        Tau = Tau,
        Σ = Σ,
        μ0 = μ0,
        Tau0 = Tau0,
        K = K,
        states = states,
        variances = variances,
        lls = lls,
        states_sm = states_sm,
        variances_sm = variances_sm,
        Covs_sm = Covs_sm,
        Y = Y,
    )

    dpr(Zs) = dpr_base(m, Zs)
    dt(Zs, Zs_new, t) = dt_base(m, Zs, Zs_new, t)
    dm(Zs, t) = dm_base(m, Zs, t)

    rinit(rng) = rinit_base(m, rng)
    dinit(Zs) = dinit_base(m, Zs)
    rp(rng, Zs, t) = rp_base(m, rng, Zs, t)
    dp(Zs, Zs_new, t) = dp_base(m, Zs, Zs_new, t)

    dpre(Zs, t) = dpre_base(m, Zs, t)

    return @NT(rinit, rp, dinit, dpr, dp, dt, dm, dpre)
end


import Base.vec

struct NominalState{T<:Real}
    p::Vector{T}
    v::Vector{T}
    q::Quaternion{T}
    ab::Vector{T}
    ωb::Vector{T}
    g::Vector{T}
end

function NominalState(x::AbstractVector)
    NominalState(
        x[1:3],   # position
        x[4:6],   # velocity
        Quaternion(x[7], x[8], x[9], x[10]), # attitude
        x[11:13], # accelerometer bias
        x[14:16], # gyrometer bias
        x[17:19]  # gravity vector
    )
end

function vec(nominal::NominalState)::AbstractVector
    [
        nominal.p
        nominal.v
        real(nominal.q)
        Quaternions.imag(nominal.q)
        nominal.ab
        nominal.ωb
        nominal.g
    ]
end

struct ErrorState{T<:Real}
    δp::Vector{T}
    δv::Vector{T}
    δθ::Vector{T}
    δab::Vector{T}
    δωb::Vector{T}
    δg::Vector{T}
end

function ErrorState(x::AbstractVector)
    ErrorState(
        x[1:3],   # position error
        x[4:6],   # velocity error
        x[7:9],   # angles error
        x[10:12], # accelerometer bias error
        x[13:15], # gyrometer bias error
        x[16:18]  # gravity error
    )
end

function vec(error::ErrorState)::AbstractVector
    [
        error.δp
        error.δv
        error.δθ
        error.δab
        error.δωb
        error.δg
    ]
end

struct Measurement{T<:Real}
    am::Vector{T}
    ωm::Vector{T}
end

struct Noise{T<:Real}
    σan::T # velocity noise Vi (m^2/s^2)
    σωn::T # orientation noise Θi (rad^2)
    σaw::T # accelero bias noise Ai (m^2/s^4)
    σωw::T # gyro bias noise Ωi (rad^2/s^2)
end


AdjMat = Matrix{Bool}
NeighborList = Vector{Vector{Int}}

function n_nodes(graph::NeighborList)
    return length(graph)
end

function n_nodes(graph::AdjMat)
    return size(graph,1)
end

function get_neighbors(graph::NeighborList, idx::Int)::Set{Int}
    # assume graph is represented in list format
    neighb = Set(graph[idx])
    return push!(neighb,idx)
end

function get_neighbors(graph::AdjMat, idx::Int)::Set{Int}
    # assume graph is represented in Adjacent matrix format
    neighb = Set([i for (i, islink) in enumerate(graph[idx,:]) if islink])
    return push!(neighb,idx)
end

function get_connected(graph, idx)::Set{Int}
    connected = Set([idx])
    to_visit = Stack{Int}()
    push!(to_visit,idx)
    while ~ isempty(to_visit)
        i = pop!(to_visit)
        push!(connected, i)
        neighb = get_neighbors(graph, i)
        for k in neighb
            if ! (k in connected)
            push!(to_visit, k)
            end
        end
    end
    return connected
end

function connected_components(graph)
    components = []
    visited = Set{Int}([])
    for i in 1:n_nodes(graph)
        if (i in visited)
            continue
        end
        comp = get_connected(graph, i)
        push!(components, comp)
        union!(visited, comp)
    end
    return components
end


include("sn.jl")
include("sn_bigr.jl")

#r=10.
r=100.
#r=110.
#r=1.1
#r=0.01
l_max = 20
n_max = l_max^2+2*l_max
sn_big = zeros(BigFloat,n_max+1);
snew = zeros(typeof(r),n_max+1);
snew_big = zeros(BigFloat,n_max+1);
#b=sqrt((r+.5)^2+0.1^2)
#b=1e-3
#b=0.99
#b=0.9
#b=0.55
b = 1+r-1e-8
#b = 1-r+1e-8
#b = 110.99
#b = 1.0
#b=100.
#b = r+1-1e-18
diff_frac = zeros(n_max+1)
diff_rel  = zeros(n_max+1)
for i=1:length(b)
  s_n_bigr!(l_max,r,b[i],snew)
  s_n_bigr!(l_max,big(r),big(b[i]),snew_big)
  s_n!(l_max,big(r),big(b[i]),sn_big)
  l = 0; n = 0; m = 0
  for n=0:n_max
    mu = l-m; nu = l+m
    if snew[n+1] != 0.0
      diff = abs(snew[n+1]/convert(Float64,snew_big[n+1])-1.)
      diff_frac[n+1]=diff
      diff_rel[n+1]=abs(snew[n+1]-convert(Float64,snew_big[n+1]))
#      if diff > 0.01
      println("n: ",n," l: ",l," m: ",m," mu: ",mu," nu: ",nu," s: ",snew[n+1],
#       " s_old: ",convert(Float64,sn_big[n+1])," d_old: ",snew[n+1]/convert(Float64,sn_big[n+1])-1.,
       " s_old: ",convert(Float64,sn_big[n+1]),
       " s_big: ",convert(Float64,snew_big[n+1])," d_old: ",snew[n+1]/convert(Float64,sn_big[n+1])-1.,
       " d_big: ",snew[n+1]/convert(Float64,snew_big[n+1])-1.)
#      end
    end
    m +=1
    if m > l
      l += 1
      m = -l
    end
  end
end

using PyPlot
clf()
semilogy(diff_frac,"o",label="fractional",c="blue")
semilogy(diff_rel,"o",label="relative",c="orange")
legend(loc="upper left")
xlabel("n")
ylabel("Error")
axis([-5,n_max,1e-25,1e-2])


# MIT License

# Copyright (c) Microsoft Corporation.

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE

"""
    RigidBodyTransform{S<:Real}

Transform encapsulating rotation and translation in 3D space. Translation happens **after** rotation.

```julia
RigidBodyTransform{S}(θ::T, ϕ::T, ψ::T, x::T, y::T, z::T)
RigidBodyTransform(rotation::SMatrix{3,3,S}, translation::SVector{3,S})
RigidBodyTransform(rotation::AbstractArray{S,2}, translation::AbstractArray{S,1})
```
`θ`, `ϕ` and `ψ` in first constructor are in **radians**.
"""
struct RigidBodyTransform{T<:Real}
    rotation::SMatrix{3,3,T,9}
    translation::SVector{3,T}

    RigidBodyTransform{S}(θ::T, ϕ::T, ψ::T, x::T, y::T, z::T) where {T<:Number,S<:Real} = new{S}(rotmat(S, θ, ϕ, ψ), SVector{3,S}(x, y, z))
    RigidBodyTransform(rotation::SMatrix{3,3,T}, translation::SVector{3,T}) where {T<:Real} = new{T}(rotation, translation)
    function RigidBodyTransform(rotation::AbstractArray{T,2}, translation::AbstractArray{T,1}) where {T<:Real}
        @assert size(rotation)[1] == size(rotation)[2] == length(translation) == 3
        return new{T}(SMatrix{3,3,T}(rotation), SVector{3,T}(translation))
    end
end
export RigidBodyTransform

"""
    identitytransform([S::Type]) -> RigidBodyTransform{S}

Returns the [`RigidBodyTransform`](@ref) of type `S` (default `Float64`) representing the identity transform.
"""
identitytransform(::Type{T} = Float64) where {T<:Real} = RigidBodyTransform(SMatrix{3,3,T,9}(I), zeros(SVector{3,T}))
export identitytransform

"""
    rotmatbetween([S::Type], a::SVector{3,T}, b::SVector{3,T}) -> SMatrix{3,3,S}

Returns the rotation matrix of type `S` (default `Float64`) representing the rotation between vetors `a` and `b`.
"""
rotmatbetween(a::SVector{3,T}, b::SVector{3,T}) where {T<:Number} = rotmatbetween(Float64, a, b)
function rotmatbetween(::Type{S}, a::SVector{3,T}, b::SVector{3,T}) where {T<:Number,S<:Real}
    v = cross(a, b)
    c = dot(a, b)
    V = SMatrix{3,3,S,9}(0, v[3], -v[2], -v[3], 0, v[1], v[2], -v[1], 0)
    R = I + V + V^2 * one(T) / (one(T) + c)
    return SMatrix{3,3,S,9}(R)
end

"""
    rotmatd([S::Type], θ::T, ϕ::T, ψ::T) -> SMatrix{3,3,S}

Returns the rotation matrix of type `S` (default `Float64`) representing the rotation by `θ`, `ϕ` and `ψ` around the *x*, *y* and *z* axes respectively **in degrees**.
"""
rotmatd(θ::T, ϕ::T, ψ::T) where {T<:Number} = rotmat(Float64, θ * π / 180, ϕ * π / 180, ψ * π / 180)
rotmatd(::Type{S}, θ::T, ϕ::T, ψ::T) where {T<:Number,S<:Real} = rotmat(S, θ * π / 180, ϕ * π / 180, ψ * π / 180)
"""
    rotmat([S::Type], θ::T, ϕ::T, ψ::T) -> SMatrix{3,3,S}

Returns the rotation matrix of type `S` (default `Float64`) representing the rotation by `θ`, `ϕ` and `ψ` around the *x*, *y* and *z* axes respectively **in radians**.
"""
rotmat(θ::T, ϕ::T, ψ::T) where {T<:Number} = rotmat(Float64, θ, ϕ, ψ)
function rotmat(::Type{S}, θ::T, ϕ::T, ψ::T) where {T<:Number,S<:Real}
    sinψ = sin(ψ)
    sinϕ = sin(ϕ)
    sinθ = sin(θ)
    cosψ = cos(ψ)
    cosϕ = cos(ϕ)
    cosθ = cos(θ)
    return SMatrix{3,3,S,9}(cosψ * cosϕ, sinψ * cosϕ, -sinϕ, cosψ * sinϕ * sinθ - sinψ * cosθ, sinψ * sinϕ * sinθ + cosψ * cosθ, cosϕ * sinθ, cosψ * sinϕ * cosθ + sinψ * sinθ, sinψ * sinϕ * cosθ - cosψ * sinθ, cosϕ * cosθ)
end
"""
    translation([S::Type], x::T, y::T, z::T) -> RigidBodyTransform{S}

Returns the [`RigidBodyTransform`](@ref) of type `S` (default `Float64`) representing the translation by `x`, `y` and `z`.
"""
translation(x::T, y::T, z::T) where {T<:Number} = translation(Float64, x, y, z)
translation(::Type{S}, x::T, y::T, z::T) where {T<:Number,S<:Real} = RigidBodyTransform(SMatrix{3,3,S,9}(I), SVector{3,S}(x, y, z))
"""
    rotation([S::Type], θ::T, ϕ::T, ψ::T) -> RigidBodyTransform{S}

Returns the [`RigidBodyTransform`](@ref) of type `S` (default `Float64`) representing the rotation by `θ`, `ϕ` and `ψ` around the *x*, *y* and *z* axes respectively **in radians**.
"""
rotation(θ::T, ϕ::T, ψ::T) where {T<:Number} = rotation(Float64, θ, ϕ, ψ)
rotation(::Type{S}, θ::T, ϕ::T, ψ::T) where {T<:Number,S<:Real} = RigidBodyTransform(rotmat(S, θ, ϕ, ψ), zeros(SVector{3,S}))
"""
    rotationd([S::Type], θ::T, ϕ::T, ψ::T) -> RigidBodyTransform{S}

Returns the [`RigidBodyTransform`](@ref) of type `S` (default `Float64`) representing the rotation by `θ`, `ϕ` and `ψ` around the *x*, *y* and *z* axes respectively **in degrees**.
"""
rotationd(θ::T, ϕ::T, ψ::T) where {T<:Number} = rotationd(Float64, θ, ϕ, ψ)
rotationd(::Type{S}, θ::T, ϕ::T, ψ::T) where {T<:Number,S<:Real} = RigidBodyTransform(rotmatd(S, θ, ϕ, ψ), zeros(SVector{3,S}))
export translation, rotation, rotationd

function Base.:*(a::RigidBodyTransform{T}, r::Ray{T,3})::Ray{T,3} where {T}
    return Ray(a * origin(r), rotate(a, direction(r)))
end

function Base.:*(a::RigidBodyTransform{T}, intsct::Intersection{T,3})::Intersection{T,3} where {T<:Real}
    u, v = uv(intsct)
    i = interface(intsct)
    if VERSION < v"1.6.0-DEV"
        # TODO REMOVE
        return @unionsplit OpticalInterface T i Intersection(α(intsct), a * point(intsct), rotate(a, normal(intsct)), u, v, i, flippednormal = flippednormal(intsct))
    else
        return Intersection(α(intsct), a * point(intsct), rotate(a, normal(intsct)), u, v, interface(intsct), flippednormal = flippednormal(intsct))
    end
end

function Base.:*(transformation::RigidBodyTransform{T}, a::Interval{T}) where {T<:Real}
    # looks ridiculous but necessary to dissambiguate the elements of the interval
    u = upper(a)
    l = lower(a)
    if l isa RayOrigin{T}
        if u isa Infinity{T}
            return Interval(l, u)
        else
            u = transformation * u
            return Interval(l, u)
        end
    else
        l = transformation * l
        if u isa Infinity{T}
            return Interval(l, u)
        else
            u = transformation * u
            return Interval(l, u)
        end
    end
end

function Base.:*(a::RigidBodyTransform{T}, tmesh::TriangleMesh{T})::TriangleMesh{T} where {T<:Real}
    newT = Vector{Triangle{T}}(undef, length(tmesh.triangles))
    @inbounds @simd for i in 1:length(tmesh.triangles)
        newT[i] = a * tmesh.triangles[i]
    end
    return TriangleMesh(newT)
end

Base.:*(a::RigidBodyTransform{T}, t::Triangle{T}) where {T<:Real} = Triangle(a * vertex(t, 1), a * vertex(t, 2), a * vertex(t, 3))

function Base.:*(a::RigidBodyTransform{T}, vector::SVector{3,T}) where {T<:Real}
    if a.rotation === SMatrix{3,3,T,9}(I)
        return vector + a.translation
    else
        # in julia 0 * Inf = NaN, so if the vector has any infinite value in it then the rotation breaks
        # have to manually test and correct this - this can still fail when summing Infs of different signs (e.g. Inf + -Inf = NaN)
        # this happens when the rotation isn't axis aligned so the result becomes less meaningful anyway
        # in this case we fall back to infinite bounding boxes anyway which are usually clipped later in the csg
        # TODO better way to do this?
        if any(isinf.(vector))
            x = zero(T)
            y = zero(T)
            z = zero(T)
            @inbounds begin
                x += abs(a.rotation[1, 1]) <= eps(T) ? zero(T) : a.rotation[1, 1] * vector[1]
                x += abs(a.rotation[1, 2]) <= eps(T) ? zero(T) : a.rotation[1, 2] * vector[2]
                x += abs(a.rotation[1, 3]) <= eps(T) ? zero(T) : a.rotation[1, 3] * vector[3]
                y += abs(a.rotation[2, 1]) <= eps(T) ? zero(T) : a.rotation[2, 1] * vector[1]
                y += abs(a.rotation[2, 2]) <= eps(T) ? zero(T) : a.rotation[2, 2] * vector[2]
                y += abs(a.rotation[2, 3]) <= eps(T) ? zero(T) : a.rotation[2, 3] * vector[3]
                z += abs(a.rotation[3, 1]) <= eps(T) ? zero(T) : a.rotation[3, 1] * vector[1]
                z += abs(a.rotation[3, 2]) <= eps(T) ? zero(T) : a.rotation[3, 2] * vector[2]
                z += abs(a.rotation[3, 3]) <= eps(T) ? zero(T) : a.rotation[3, 3] * vector[3]
            end
            return SVector(x, y, z) + a.translation
        else
            return a.rotation * vector + a.translation
        end
    end
end
Base.:*(a::RigidBodyTransform{T}, b::RigidBodyTransform{T}) where {T<:Real} = RigidBodyTransform(a.rotation * b.rotation, a.translation + a.rotation * b.translation)
rotate(a::RigidBodyTransform{T}, vector::SVector{3,T}) where {T<:Real} = a.rotation * vector
Base.inv(a::RigidBodyTransform{T}) where {T<:Real} = RigidBodyTransform(a.rotation', -a.rotation' * a.translation)
Base.print(io::IO, a::RigidBodyTransform) = print(io, hcat(a.rotation, a.translation))
Base.collect(a::RigidBodyTransform) = hcat(a.rotation, a.translation)