#Demo 13: Broadcast test: Slicing a Paraboloid
#-------------------------------------------------------------------------------

using CMDimCircuits
CMDimCircuits.@using_CData()

#Get a demo display:
include(CMDimCircuits.demoplotcfgscript)
#Normally use something like:
#CMDimData.@includepkg EasyPlotInspect

#==Constants
===============================================================================#
a_yvsx = cons(:a, labels = set(yaxis="Y", xaxis="X"))
dfltwfrmattr = cons(:a,
	line = set(style=:none, width=2),
	glyph = set(shape=:o),
)

#==Input data
===============================================================================#
nsteps = 100 #Number of steps to sample x/y values of paraboloid
pext = 10 #sample paraboloid up to x/y = +/-pext

#Paraboloid coefficients
a = 1
b = 2
c = 1

#Axis ranges based on input data:
axrange = cons(:a, xyaxes=set(xmin=-pext, xmax=pext, ymin=-pext, ymax=pext))

#==Computations
===============================================================================#
paraboloid(x,y) = ((x/a)^2 + (y/b)^2) * c

#Generate parameter sweeps:
#NOTE: x is outermost sweep (because outermost sweep plots on x-axis)
sweeplist = PSweep[
	PSweep("y", collect(-pext:(2pext/nsteps):pext))
	PSweep("x", collect(-pext:(2pext/nsteps):pext))
]

para = DataHR{DataFloat}(sweeplist) #Create empty pattern
plane1 = DataHR{DataFloat}(sweeplist) #Create empty pattern
plane2 = DataHR{DataFloat}(sweeplist) #Create empty pattern
for inds in subscripts(para)
	(y,x) = coordinates(para, inds)
	para.elem[inds...] = paraboloid(x,y)
	plane1.elem[inds...] = 3x+y/4+2
	plane2.elem[inds...] = -5(x+10)+5y+50
end

slice_z3 = xcross(para-3) #Cross @ z=3: (x-values@x-coord)
slice_z60 = xcross(para-60) #Cross @ z=60
slice_p1 = xcross(para-plane1) #Cross @ plane1
slice_p2 = xcross(para-plane2) #Cross @ plane1

#Keep x-coordinate, but translate up by y-coordinate:
for dsym in (:slice_z3, :slice_z60, :slice_p1, :slice_p2)
@eval begin
_y = parameter($dsym, "y") #Really only needs to be computed once
$dsym = $dsym*0+_y #Keep x-coord, but translate up by y-coordinate of $dsym
end
end

#==Generate plot
===============================================================================#
p1 = push!(cons(:plot, a_yvsx, axrange, title="Slice with planes z=3 & z=60", legend=false),
	cons(:wfrm, slice_z3, dfltwfrmattr, label="z=3"),
	cons(:wfrm, slice_z60, dfltwfrmattr, label="z=60"),
)
p2 = push!(cons(:plot, a_yvsx, axrange, title="Slice with planes 1 & 2", legend=false),
	cons(:wfrm, slice_p1, dfltwfrmattr, label="plane 1"),
	cons(:wfrm, slice_p2, dfltwfrmattr, label="plane 1"),
)

pcoll = push!(cons(:plotcoll, title="Sliced Paraboloid"), p1, p2)
	pcoll.theme.colorscheme = EasyPlot.ColorScheme(EasyPlot.colormap("Blues", 12))
	pcoll.ncolumns = 1

#==Display results in pcoll
===============================================================================#
EasyPlot.displaygui(pcoll)

#==Return pcoll to user (call evalfile(...))
===============================================================================#
pcoll


abstract type Separator end

"""
ConstraintSeparator is a separator that represents a constraint defined directly
using `@constraint`.
"""
struct ConstraintSeparator{C, II, ex<:Union{Operation,Expr}} <: Separator
    variables::Vector{Symbol}
    constraint::II  # Interval or IntervalBox
    contractor::C
    expression::ex
end

ConstraintSeparator(constraint, contractor, expression) = ConstraintSeparator(contractor.variables, constraint, contractor, expression)

"""CombinationSeparator is a separator that is a combination (union, intersection,
or complement) of other separators.
"""
struct CombinationSeparator{F, ex<:Union{Operation,Expr}} <: Separator
    variables::Vector{Symbol}
    separator::F
    expression::ex
end

function (S::ConstraintSeparator)(X::IntervalBox)
    C = S.contractor
    a, b = S.constraint.lo, S.constraint.hi

    inner = C(IntervalBox(Interval(a, b)), X)

    local outer

    if a == -∞
        outer = C(IntervalBox(Interval(b, ∞)), X)

    elseif b == ∞
        outer = C(IntervalBox(Interval(-∞, a)), X)

    else

        outer1 = C(IntervalBox(Interval(-∞, a)), X)
        outer2 = C(IntervalBox(Interval(b, ∞)), X)

        outer = outer1 ∪ outer2
    end

    return (inner, outer)
end

"""`parse_comparison` parses comparisons like `x >= 10`
into the corresponding interval, expressed as `x ∈ [10,∞]`

Returns the expression and the constraint interval

TODO: Allow something like [3,4]' for the complement of [3,4]
"""

function parse_comparison(ex::Expr)
    expr, limits =
    @match ex begin
       ((a_ <= b_) | (a_ < b_) | (a_ ≤ b_))   => (a, (-∞, b))
       ((a_ >= b_) | (a_ > b_) | (a_ ≥ b_))   => (a, (b, ∞))
       ((a_ == b_) | (a_ = b_))   => (a, (b, b))

       ((a_ <= b_ <= c_)
        | (a_ < b_ < c_)
        | (a_ <= b_ < c)
        | (a_ < b_ <= c))         => (b, (a, c))

       ((a_ >= b_ >= c_)
       | (a_ > b_ > c_)
       | (a_ >= b_ > c_)
       | (a_ > b_ >= c))          => (b, (c, a))

       ((a_ ∈ [b_, c_])
       | (a_ in [b_, c_])
       | (a_ ∈ b_ .. c_)
       | (a_ in b_ .. c_))        => (a, (b, c))

       _                          => (ex, (-∞, ∞))

   end

   a, b = limits
   return (expr, a..b)   # expr ∈ [a,b]

end

function parse_comparison(ex::Operation)

    if isa(ex.args[1], ModelingToolkit.Constant)
        if ex.op == <
            a = ex.args[1].value
            b = Inf
        elseif ex.op == >
            a = -Inf
            b = ex.args[1].value
        end
        return (ex.args[2], a..b)
    elseif isa(ex.args[2], ModelingToolkit.Constant)
        if ex.op == >
            a = ex.args[2].value
            b = Inf
        elseif ex.op == <
            a = -Inf
            b = ex.args[2].value
        else
            a = ex.args[2].value
            b = ex.args[2].value
        end
        return (ex.args[1], a..b)
    end

end

function new_parse_comparison(ex)
    # @show ex
    if @capture ex begin
            (op_(a_, b_))
        end

        #return (op, a, b)
        # @show op, a, b

    elseif ex.head == :comparison
        println("Comparison")
        symbols = ex.args[1:2:5]
        operators = ex.args[2:2:4]

        # @show symbols
        # @show operators

    end
end

function make_constraint(expr, constraint, var =[])

    if isa(expr, Symbol)
        expr = :(1 * $expr)  # make into an expression!
    end

    contractor_name = make_symbol(:C)

    full_expr = Meta.quot(:($expr ∈ $constraint))

    contractor_code = make_contractor(expr, var)

    code = quote end
    push!(code.args, :($(esc(contractor_name)) = $(contractor_code)))

    push!(code.args, :(ConstraintSeparator($constraint, $(esc(contractor_name)), $full_expr)))

    code
end

make_constraint(expr::Variable, constraint) = make_constraint(Operation(expr), constraint)

function make_constraint(expr::Operation, constraint, var=[])
    C = Contractor(var, expr)
    ex = expr ∈ constraint
    ConstraintSeparator(constraint, C, ex)
end

"""Create a separator from a given constraint expression, written as
standard Julia code.

e.g. `C = @constraint x^2 + y^2 <= 1`

The variables (`x` and `y`, in this case) are automatically inferred.
External constants can be used as e.g. `\$a`:

```
a = 3
C = @constraint x^2 + y^2 <= \$a
```
"""
macro constraint(ex::Expr, variables = [])
    expr, constraint = parse_comparison(ex)
    isa(variables, Array) ? var = [] : var = variables.args
    make_constraint(expr, constraint, var)
end

"""
Create a separator without the use of macros using ModelingToolkit

e.g  vars = @variables x y z
S = Separator(vars, x^2+y^2<1)
X= IntervalBox(-0.5..1.5, -0.5..1.5, -0.5..1.5)
S(X)
"""
function Separator(variables, ex::Operation)
    expr, constraint = parse_comparison(ex)
    make_constraint(expr, constraint, variables)
end

Separator(ex::Operation) = Separator([], ex)

Separator(vars::Union{Vector{Operation}, Tuple{Vararg{Operation,N}}}, g::Function) where N = Separator(vars, g(vars...))

Separator(vars, f::Function) = Separator([Variable(Symbol(i))() for i in vars], f([Variable(Symbol(i))() for i in vars]...)) # if vars is not vector of variables

function show(io::IO, S::Separator)
    println(io, "Separator:")
    print(io, "  - variables: ")
    print(io, join(map(string, S.variables), ", "))
    println(io)
    print(io, "  - expression: ")
    print(io, S.expression)
end

(S::CombinationSeparator)(X) = S.separator(X)

"Unify the variables of two separators"
function unify_variables(vars1, vars2)

    variables = unique(sort(vcat(vars1, vars2)))
    numvars = length(variables)  # total number of variables

    indices1 = indexin(vars1, variables)
    indices2 = indexin(vars2, variables)

    where1 = zeros(Int, numvars)  # inverse of indices1
    where2 = zeros(Int, numvars)

    for (i, which) in enumerate(indices1)
        where1[which] = i
    end

    for (i, which) in enumerate(indices2)
        where2[which] = i
    end

    return variables, indices1, indices2, where1, where2
end

# TODO: when S1 and S2 have different variables -- amalgamate!

"""
    ∩(S1::Separator, S2::Separator)

Separator for the intersection of two sets given by the separators `S1` and `S2`.
Takes an iterator of intervals (`IntervalBox`, tuple, array, etc.), of length
equal to the total number of variables in `S1` and `S2`;
it returns inner and outer tuples of the same length
"""
function ∩(S1::Separator, S2::Separator)

    #=variables, indices1, indices2, where1, where2 = unify_variables(S1.variables, S2.variables)
    numvars = length(variables)=#
    variables = S1.variables   # as S1 and S2 have same variables
    f = X -> begin

       inner1, outer1 = S1(X)
       inner2, outer2 = S2(X)
       inner = inner1 ∩ inner2
       outer = outer1 ∪ outer2

        #=
        inner1, outer1 = S1(IntervalBox([X[i] for i in indices1]...))
        inner2, outer2 = S2(IntervalBox([X[i] for i in indices2]...))

        if any(isempty, inner1)
            inner1 = emptyinterval(IntervalBox(X))
        else
            inner1 = [i ∈ indices1 ? inner1[where1[i]] : X[i] for i in 1:numvars]
        end

        if any(isempty, outer1)
            outer1 = emptyinterval(IntervalBox(X))
        else
            outer1 = [i ∈ indices1 ? outer1[where1[i]] : X[i] for i in 1:numvars]
        end

        if any(isempty, inner2)
            inner2 = emptyinterval(IntervalBox(X))
        else
            inner2 = [i ∈ indices2 ? inner2[where2[i]] : X[i] for i in 1:numvars]
        end

        if any(isempty, outer2)
            outer2 = emptyinterval(IntervalBox(X))
        else
            outer2 = [i ∈ indices2 ? outer2[where2[i]] : X[i] for i in 1:numvars]
        end

        # Treat as if had X[i] in the other directions, except if empty

        inner = IntervalBox( [x ∩ y for (x,y) in zip(inner1, inner2) ]... )
        outer = IntervalBox( [x ∪ y for (x,y) in zip(outer1, outer2) ]... )
        =#

        return (inner, outer)

    end

    expression = :($(S1.expression) ∩ $(S2.expression))

    return CombinationSeparator(variables, f, expression)

end

function ∪(S1::Separator, S2::Separator)

    #=variables, indices1, indices2, where1, where2 = unify_variables(S1.variables, S2.variables)
    numvars = length(variables)=#
    variables = S1.variables    # S1 and S2 have same variables
    f = X -> begin

        inner1, outer1 = S1(X)
        inner2, outer2 = S2(X)
        inner = inner1 ∪ inner2
        outer = outer1 ∩ outer2

        #=
        inner1, outer1 = S1(IntervalBox([X[i] for i in indices1]...))
        inner2, outer2 = S2(IntervalBox([X[i] for i in indices2]...))

        if any(isempty, inner1)
            inner1 = emptyinterval(IntervalBox(X))
        else
            inner1 = [i ∈ indices1 ? inner1[where1[i]] : X[i] for i in 1:numvars]
        end

        if any(isempty, outer1)
            outer1 = emptyinterval(IntervalBox(X))
        else
            outer1 = [i ∈ indices1 ? outer1[where1[i]] : X[i] for i in 1:numvars]
        end

        if any(isempty, inner2)
            inner2 = emptyinterval(IntervalBox(X))
        else
            inner2 = [i ∈ indices2 ? inner2[where2[i]] : X[i] for i in 1:numvars]
        end

        if any(isempty, outer2)
            outer2 = emptyinterval(IntervalBox(X))
        else
            outer2 = [i ∈ indices2 ? outer2[where2[i]] : X[i] for i in 1:numvars]
        end

        inner = IntervalBox( [x ∪ y for (x,y) in zip(inner1, inner2) ]... )
        outer = IntervalBox( [x ∩ y for (x,y) in zip(outer1, outer2) ]... )
        =#

        return (inner, outer)
    end

    expression = :($(S1.expression) ∪ $(S2.expression))

    return CombinationSeparator(variables, f, expression)

end

function !(S::Separator)
    f = X -> begin
        inner, outer = S(X)
        return (outer, inner)
    end

    expression = :(!($(S.expression)))

    return CombinationSeparator(S.variables, f, expression)
end


## script for Fig.7, Fig.8(a), Fig.10(a)
using VoronoiDelaunay, VoronoiCells, GeometricalPredicates
using MultiscaleGraphSignalTransforms, Plots, LightGraphs, JLD; gr(dpi = 200)

barbara = JLD.load(joinpath(@__DIR__, "..", "datasets", "barbara_gray_matrix.jld"), "barbara")
G, L, X = SunFlowerGraph(); N = nv(G)
#################### Fig. 7(a) sunflower graph
gplot(1.0*adjacency_matrix(G),X; width=1); scatter_gplot!(X; c = :red, ms = LinRange(1,9,N)); plt = plot!(frame = :none)
# savefig(plt, joinpath(@__DIR__, "../paperfigs/SunFlower.png"))

## Voronoi tessellation
width_x = maximum(abs.(X[:,1])) * 2; width_y = maximum(abs.(X[:,2])) * 2;
width = VoronoiDelaunay.max_coord - VoronoiDelaunay.min_coord
center_coord = (VoronoiDelaunay.min_coord + VoronoiDelaunay.max_coord)/2
X_transform = zeros(N,2)
for i in 1:N
    X_transform[i,:] = X[i,:] ./ [width_x/width, width_y/width] + [center_coord, center_coord]
end

pts = [Point2D(X_transform[i,1], X_transform[i,2]) for i in 1:N]
tess = DelaunayTessellation(N)
push!(tess, pts)
#################### Fig. 7(b) voronoi tessellation
xx, yy = getplotxy(voronoiedges(tess))
plt = plot(xx, yy, xlim=[1,2], ylim=[1,2], linestyle=:auto, linewidth=1, linecolor=:blue, grid=false, label="", aspect_ratio=1, frame=:box)
# savefig(plt, joinpath(@__DIR__, "../paperfigs/Sunflower_Barbara_Voronoi_cells.png"))

#################### Fig. 8(a) barbara sunflower eye
heatmap(barbara, yflip=true, ratio=1, c=:greys); scatter_gplot!(transform2D(X; s = 20, t = [395, 100]); ms = 2, c = :red); sample_location_plt = plot!(cbar = false, frame = :none)
# savefig(sample_location_plt, joinpath(@__DIR__, "../paperfigs/barb_sunflower_eye.png"))

#################### Fig. 10(a) barbara sunflower pants
heatmap(barbara, yflip=true, ratio=1, c=:greys); scatter_gplot!(transform2D(X; s = 20, t = [280, 320]); ms = 2, c = :red); sample_location_plt = plot!(cbar = false, frame = :none)
# savefig(sample_location_plt, joinpath(@__DIR__, "../paperfigs/barb_sunflower_pants.png"))


using Printf
using GLMNet: glmnet

"""
    screen_lasso(X, y, pmax)

Screening out `pmax` variables by lasso (glmnet) on `X` and `y`.
"""
function screen_lasso(X::Matrix{Float64}, y::Vector{Float64}, pmax::Int64)
    fit = glmnet(X, y)
    p = size(X, 2)
    @assert pmax <= p
    betas = fit.betas
    n_vars = sum(betas.!=0, dims=1)[:]
    S = Set{Int64}()
    for n in unique(n_vars)
        Z = betas[:, n_vars.==n]
        non_zeros = sum(Z.!=0, dims=2)[:]
        vars = (1:p)[non_zeros .> 0]
        new_vars = setdiff(vars, S)
        Z = Z[new_vars, :]
        if length(S) + length(new_vars) > pmax
            # must break ties
            z = abs.(Z[:, end])
            idx = sort(1:length(z), by = i -> z[i], rev = true)
            new_vars =  new_vars[idx[1:(pmax - length(S))]]
        end
        union!(S, new_vars)
        length(S) == pmax && break
    end
    return sort(collect(S))
end

"""
    screen_HOLP(X, y, pmax)

Screen out `pmax` variables with HOLP projection.

See: Wang, Xiangyu, and Chenlei Leng. "High dimensional ordinary least squares projection for screening variables."
     Journal of the Royal Statistical Society: Series B (Statistical Methodology) 78.3 (2016): 589-611.

"""
function screen_HOLP(X::Matrix{Float64}, y::Vector{Float64}, pmax::Int64)
    n, p = size(X)
    @assert p >= n
    @assert pmax < p
    _X = X .- mean(X, 1)
    _X = _X ./ std(_X, 1)
    _y = y - mean(y)
    β = _X' * ((_X * _X' + 10 * eye(n)) \ _y)
    return sortperm(abs.(β), rev=true)[1:pmax]
end


import Base.*
import Base.size
import Base.eltype
import LinearAlgebra.issymmetric

import LinearAlgebra.mul!

struct SqLinOp{Tv,Ti}
    issym::Bool
    value::Tv
    n::Ti
    multFn::Function #this is bad in terms of types?

    SqLinOp{Tv,Ti}(issym,value,n,multFn) where {Tv,Ti} = new(issym,value,n,multFn)
end

SqLinOp(issym,value::Tv,index::Ti,multFn) where {Tv,Ti} = SqLinOp{Tv,Ti}(issym,value,index,multFn)

eltype(A::SqLinOp{Tv,Ti}) where {Tv,Ti} = Tv

size(A::SqLinOp{Tv,Ti}, d::Ti) where {Tv,Ti} = A.n

size(A::SqLinOp{Tv,Ti}) where {Tv,Ti} = (A.n,A.n)

issymmetric(A::SqLinOp{Tv,Ti}) where {Tv,Ti} = A.issym

function *(A::SqLinOp{Tv,Ti}, b) where {Tv,Ti}
    return A.multFn(b)
end

# put in for compat with Arpack in Julia v0.7
function mul!(Y, A::SqLinOp{Tv,Ti}, B) where {Tv,Ti}
    Y1 = A*B
    for i in 1:A.n
        Y[i] = Y1[i]
    end
end

function testId(n::Ti) where Ti
    return M = SqLinOp{Float64,Int64}(true,1.0,n,x -> x)
end


function characteristic_matrix_ellipse(q0::Float64, q1::Float64, q2::Float64, q3::Float64, r1::Float64, r2::Float64)

	a11::Float64 = 4.0 * r2^2 * (q0 * q3 + q1 * q2)^2 + r1^2 * (q0^2 + q1^2 - q2^2 - q3^2)^2
	a12::Float64 = r2^2 * (2.0 * q0 * q3 + 2.0 * q1 * q2) * (q0^2 - q1^2 + q2^2 - q3^2) - r1^2 * (2.0 * q0 * q3 - 2.0 * q1 * q2) * (q0^2 + q1^2 - q2^2 - q3^2)
	a13::Float64 = r1^2 * (2.0 * q0 * q2 + 2.0 * q1 * q3) * (q0^2 + q1^2 - q2^2 - q3^2) - r2^2 * (2.0 * q0 * q1 - 2.0 * q2 * q3) * (2.0 * q0 * q3 + 2.0 * q1 * q2)
	a21::Float64 = a12
	a22::Float64 = 4.0 * r1^2 * (q0 * q3 - q1 * q2)^2 + r2^2 * (q0^2 - q1^2 + q2^2 - q3^2)^2
	a23::Float64 = - r2^2 * (2.0 * q0 * q1 - 2.0 * q2 * q3) * (q0^2 - q1^2 + q2^2 - q3^2) - r1^2 * (2.0 * q0 * q2 + 2.0 * q1 * q3) * (2.0 * q0 * q3 - 2.0 * q1 * q2)
	a31::Float64 = a13
	a32::Float64 = a23
	a33::Float64 = 4.0 * r1^2 * (q0 * q2 + q1 * q3)^2 + 4.0 * r2^2 * (q0 * q1 - q2 * q3)^2

	return (a11, a12, a13, a21, a22, a23, a31, a32, a33)
end


include("backward_linear_steps.jl")

"""
compute the gradients (∇) of the parameters (master_cache) of the constructed model with respect to the cost of prediciton (Ŷ) in comparison with actual output (Y).
"""
function back_propagate_model_weights(Ŷ, Y, master_cache)
    # initiate the dictionary to store the gradients for all the components in each layer
    ∇ = Dict()

    L = length(master_cache)
    Y = reshape(Y, size(Ŷ))

    # Partial derivate of the output layer
    δŶ = - (Y / Ŷ) + ((1 .- Y) / (1 .- Ŷ))
    current_cache = master_cache[L]

    # backpropagate on the layer preceeding the output layer
    ∇[string("δW_", (L))], ∇[string("δb_", (L))], ∇[string("δA_", (L-1))] = linear_activation_backwards(δŶ, current_cache, activation_function = "softmax")

    # go backward in the layers and compute the partial derivates of each component
    for l = reverse(0:L-2)
        current_cache = master_cache[l+1]
        ∇[string("δW_", (l+1))], ∇[string("δb_", (l+1))], ∇[string("δA_", (l))] = linear_activation_backwards(∇[string("δA_", (l+1))], current_cache, activation_function="tanh")
    end # for

    return ∇
end # function


using JuMP, EAGO

                     m = Model()

                     EAGO.register_eago_operators!(m)

                     @variable(m, -1 <= x[i=1:3] <= 1)

                     @variable(m, -5.332774675925122 <= q <= 6.4153373403467775)

                     add_NL_constraint(m, :(sigmoid(0.1742606373223392 + -0.05222015296022331*$(x[1]) + 0.15266538164806986*$(x[2]) + 0.27567618482533085*$(x[3])) + sigmoid(0.7992746561529409 + 0.8053279528303205*$(x[1]) + -0.8864483595131469*$(x[2]) + 0.6110618799366847*$(x[3])) + sigmoid(-0.6752579159956698 + 0.27352401436271423*$(x[1]) + 0.2855995340868476*$(x[2]) + -0.33283297885270313*$(x[3])) + sigmoid(0.2430039547312175 + -0.8377319978555251*$(x[1]) + 0.39266174748446403*$(x[2]) + -0.9683058237799194*$(x[3])) - $q <= 0.0))

                     @objective(m, Min, q)

                     return m


using Test

include("rational-numbers.jl")

@test RationalNumber <: Real
@test_throws ArgumentError RationalNumber(0, 0)

@testset "One- & Zero-elements" begin
    @test zero(RationalNumber{Int}) == RationalNumber(0, 1)
    @test one(RationalNumber{Int})  == RationalNumber(1, 1)
end

@testset "Arithmetic" begin
    @testset "Addition" begin
        @test RationalNumber( 1, 2) + RationalNumber( 2, 3) == RationalNumber( 7, 6)
        @test RationalNumber( 1, 2) + RationalNumber(-2, 3) == RationalNumber(-1, 6)
        @test RationalNumber(-1, 2) + RationalNumber(-2, 3) == RationalNumber(-7, 6)
        @test RationalNumber( 1, 2) + RationalNumber(-1, 2) == RationalNumber( 0, 1)
    end

    @testset "Subtraction" begin
        @test RationalNumber( 1, 2) - RationalNumber( 2, 3) == RationalNumber(-1, 6)
        @test RationalNumber( 1, 2) - RationalNumber(-2, 3) == RationalNumber( 7, 6)
        @test RationalNumber(-1, 2) - RationalNumber(-2, 3) == RationalNumber( 1, 6)
        @test RationalNumber( 1, 2) - RationalNumber( 1, 2) == RationalNumber( 0, 1)
    end

    @testset "Multiplication" begin
        @test RationalNumber( 1, 2) * RationalNumber( 2, 3) == RationalNumber( 1, 3)
        @test RationalNumber(-1, 2) * RationalNumber( 2, 3) == RationalNumber(-1, 3)
        @test RationalNumber(-1, 2) * RationalNumber(-2, 3) == RationalNumber( 1, 3)
        @test RationalNumber( 1, 2) * RationalNumber( 2, 1) == RationalNumber( 1, 1)
        @test RationalNumber( 1, 2) * RationalNumber( 1, 1) == RationalNumber( 1, 2)
        @test RationalNumber( 1, 2) * RationalNumber( 0, 1) == RationalNumber( 0, 1)
    end

    @testset "Exponentiation" begin
        @testset "Exponentiation of a rational number" begin
            @test RationalNumber( 1, 2)^3 == RationalNumber( 1, 8)
            @test RationalNumber(-1, 2)^3 == RationalNumber(-1, 8)
            @test RationalNumber( 0, 1)^5 == RationalNumber( 0, 1)
            @test RationalNumber( 1, 1)^4 == RationalNumber( 1, 1)
            @test RationalNumber( 1, 2)^0 == RationalNumber( 1, 1)
            @test RationalNumber(-1, 2)^0 == RationalNumber( 1, 1)
        end

        @testset "Exponentiation of a real number to a rational number" begin
            @test 8^RationalNumber( 4, 3) ≈ 15.999999999999998
            @test 9^RationalNumber(-1, 2) ≈ 0.3333333333333333
            @test 2^RationalNumber( 0, 1) ≈ 1.0
        end
    end

    @testset "Division" begin
        @test RationalNumber( 1, 2) / RationalNumber( 2, 3) == RationalNumber( 3, 4)
        @test RationalNumber( 1, 2) / RationalNumber(-2, 3) == RationalNumber(-3, 4)
        @test RationalNumber(-1, 2) / RationalNumber(-2, 3) == RationalNumber( 3, 4)
        @test RationalNumber( 1, 2) / RationalNumber( 1, 1) == RationalNumber( 1, 2)
    end
end

@testset "Absolute value" begin
    @test abs(RationalNumber( 1,  2)) == RationalNumber(1, 2)
    @test abs(RationalNumber(-1,  2)) == RationalNumber(1, 2)
    @test abs(RationalNumber( 0,  1)) == RationalNumber(0, 1)
    @test abs(RationalNumber( 1, -2)) == RationalNumber(1, 2)
end

@testset "Reduction to lowest terms" begin
    r = RationalNumber(2, 4)
    @test numerator(r)   == 1
    @test denominator(r) == 2

    r = RationalNumber(-4, 6)
    @test numerator(r)   == -2
    @test denominator(r) ==  3

    r = RationalNumber(3, -9)
    @test numerator(r)   == -1
    @test denominator(r) ==  3

    r = RationalNumber(0, 6)
    @test numerator(r)   == 0
    @test denominator(r) == 1

    r = RationalNumber(-14, 7)
    @test numerator(r)   == -2
    @test denominator(r) ==  1

    r = RationalNumber(13, 13)
    @test numerator(r)   == 1
    @test denominator(r) == 1

    r = RationalNumber(1, -1)
    @test numerator(r)   == -1
    @test denominator(r) ==  1
end

# TODO add to problem spec
# The following testset is based on the tests for rational numbers in Julia Base (MIT license)
# https://github.com/JuliaLang/julia/blob/52bafeb981bac548afd2264edb518d8d86944dca/test/rational.jl
# https://github.com/JuliaLang/julia/blob/52bafeb981bac548afd2264edb518d8d86944dca/LICENSE.md
@testset "Ordering" begin
    for a in -5:5, b in -5:5, c in -5:5
        a == b == 0 && continue

        r = RationalNumber(a, b)

        @test (r == c) == (a / b == c)
        @test (r != c) == (a / b != c)
        @test (r <= c) == (a / b <= c)
        @test (r <  c) == (a / b <  c)
        @test (r >= c) == (a / b >= c)
        @test (r >  c) == (a / b >  c)

        for d in -5:5
            c == d == 0 && continue

            s = RationalNumber(c, d)

            @test (r == s) == (a / b == c / d)
            @test (r != s) == (a / b != c / d)
            @test (r <= s) == (a / b <= c / d)
            @test (r <  s) == (a / b <  c / d)
            @test (r >= s) == (a / b >= c / d)
            @test (r >  s) == (a / b >  c / d)
        end
    end
end

@testset "Showing RationalNumbers" begin
    @test sprint(show, RationalNumber(23, 42)) == "23//42"
    @test sprint(show, RationalNumber(-2500, 5000)) == "-1//2"
end


using PyPlot, AcousticWave
nz = 201; nx = 201; ext = 30; iflag = 1; dz = 5.; dx = 5.; dt = 1e-3; f0 = 30.0
v  = 2500. * ones(nz, nx);
v[101:end, 1:end] = 3000.;
vmax = maximum(v); vmin = minimum(v);
fidMtx = InitFidMtx(nz, nx, ext, iflag, dz, dx, dt, vmax, vmin, f0, v);
# initialize source term
f0 = 30.0; ot = [0.0]; pos_src = [5 101]
src = InitSources(pos_src, f0, ot, dt);
# receiver location
irx = collect(1:2:nx); irz = 5*ones(Int64, length(irx));
pos_rec = convert(Array{Int64,2}, hcat(irz, irx))
# position of receiver
shot  = MultiStepForward(pos_rec, src[1], fidMtx, tmax=1.0)
SeisPlot(shot.d, pclip=95, cmap="seismic", name="/Users/wenlei/Desktop/shot.pdf")
dp = dpDirect(shot, v[1], src[1], dx, dt)
shot.d = dp .* shot.d
SeisPlot(shot.d, cmap="seismic", name="/Users/wenlei/Desktop/shot_rd.pdf")

# compute source side wavefield with smooth velocity model
v1 = modSmooth(v, 10)
fidMtx1 = InitFidMtx(nz, nx, ext, iflag, dz, dx, dt, vmax, vmin, f0, v1);
pathsrc = "/Users/wenlei/Desktop/src.bin"
MultiStepForward(pathsrc, src, fidMtx1, tmax=1.0)
pathout  = "/Users/wenlei/Desktop/srcWave"
waveAnim(pathsrc, pathout, interval=100)

# compute receiver side wavefield with smooth velocity model
pathrec = "/Users/wenlei/Desktop/rec.bin"
MultiStepAdjoint(pathrec, shot, fidMtx1)
animRec = "/Users/wenlei/Desktop/recWave"
waveAnim(pathrec, animRec, interval=100)
(Ivz, Ivx, Ipz, Ipx) = RTMimaging(pathsrc, pathrec, image_type="all");
SeisPlot(Ipz, cmap="seismic")
Ip = RTMimaging(pathsrc, pathrec, image_type="p");
SeisPlot(Ip, cmap="seismic")

# ==============================================================================
using PyPlot, AcousticWave
nz = 101; nx = 101; ext = 30; iflag = 1; dz = 5.; dx = 5.; dt = 1e-3; f0 = 30.0
v  = 2500. * ones(nz, nx);
v[101:end, 1:end] = 3000.;
vmax = maximum(v); vmin = minimum(v);
fidMtx = InitFidMtx(nz, nx, ext, iflag, dz, dx, dt, vmax, vmin, f0, v);
m31 = fidMtx.MpzBvz * fidMtx.MvzBvz;
m33 = fidMtx.MpzBvz * fidMtx.MvzBp + fidMtx.MpzBpz;
m34 = fidMtx.MpzBvz * fidMtx.MvzBp;

m42 = fidMtx.MpxBvx * fidMtx.MvxBvx;
m43 = fidMtx.MpxBvx * fidMtx.MvxBp;
m44 = fidMtx.MpxBvx * fidMtx.MvxBp + fidMtx.MpxBpx;

N = (nz+2*ext)*(nx+2*ext);
M = vcat(hcat(fidMtx.MvzBvz, spzeros(N,N) , fidMtx.MvzBp, fidMtx.MvzBp),
         hcat(spzeros(N,N) , fidMtx.MvxBvx, fidMtx.MvxBp, fidMtx.MvxBp),
         hcat(m31          , spzeros(N,N) , m33         , m34         ),
         hcat(spzeros(N,N) , m42          , m43         , m44         ))

# ==============================================================================
# using PyPlot, AcousticWave
# nz = 101; nx = 101; ext = 30; iflag = 1; dz = 5.; dx = 5.; dt = 1e-3; f0 = 30.0
# v  = 3000. * ones(nz, nx);
# v[101:end, 1:end] = 3000.;
# vmax = maximum(v); vmin = minimum(v);
# fidMtx = InitFidMtx(nz, nx, ext, iflag, dz, dx, dt, vmax, vmin, f0, v);
# fidMtx1= InitFidMtx_back(nz, nx, ext, iflag, dz, dx, dt, vmax, vmin, f0, v);
#
# function OneStepBackward(spt2::SnapShot, spt1::SnapShot, fidMtx::FidMtx)
#     spt2.it = spt1.it - 1
#     spt2.pz = fidMtx1.MpzBpz * spt1.pz + fidMtx1.MpzBvz * spt1.vz;
#     spt2.px = fidMtx1.MpxBpx * spt1.px + fidMtx1.MpxBvx * spt1.vx;
#     spt2.vz = fidMtx1.MvzBvz * spt1.vz + fidMtx1.MvzBp * (spt2.pz+spt2.px);
#     spt2.vx = fidMtx1.MvxBvx * spt1.vx + fidMtx1.MvxBp * (spt2.pz+spt2.px);
#     return nothing
# end
#
# spt1 = InitSnapShot(0, 0, nz, nx, ext, iflag, dt, 2);
# spt2 = InitSnapShot(0, 0, nz, nx, ext, iflag, dt, 2);
# spt3 = InitSnapShot(0, 0, nz, nx, ext, iflag, dt, 2);
# spt1.vz = rand(length(spt1.vz)); spt1.vx = rand(length(spt1.vx));
# spt1.pz = rand(length(spt1.pz)); spt1.px = rand(length(spt1.px));
# OneStepBackward(spt2, spt1, fidMtx);
# OneStepAdjoint!(spt3, spt1, fidMtx);
#
# p2 = spt2.pz + spt2.px
# p3 = spt3.pz + spt3.px
#
# tmp = spt3-spt2


@system Area begin
    NU: nodal_units ~ hold
    PD: planting_density ~ hold

    leaf_area(x=NU["*"].leaf.area) => sum(x) ~ track(u"cm^2")

    green_leaf_area(x=NU["*"].leaf.green_area) => sum(x) ~ track(u"cm^2")

    green_leaf_ratio(green_leaf_area, leaf_area) => begin
        iszero(leaf_area) ? 0. : (green_leaf_area / leaf_area)
    end ~ track

    LAI(green_leaf_area, PD): leaf_area_index => begin
        green_leaf_area * PD
    end ~ track

    # actualgreenArea is the green area of leaf growing under carbon limitation
    #SK 8/22/10: There appears to be no distinction between these two variables in the code.
    actual_green_leaf_area(green_leaf_area) => green_leaf_area ~ track(u"cm^2")

    senescent_leaf_area(x=NU["*"].leaf.senescent_area) => sum(x) ~ track(u"cm^2")

    potential_leaf_area(x=NU["*"].leaf.potential_area) => sum(x) ~ track(u"cm^2")

    potential_leaf_area_increase(x=NU["*"].leaf.potential_area_increase) => sum(x) ~ track(u"cm^2")

    # calculate relative area increases for leaves now that they are updated
    #TODO remove if unnecessary
    # relative_leaf_area_increase(x=NU["*"].leaf.relative_area_increase) => sum(x) ~ track(u"cm^2")

    #FIXME it doesn't seem to be 'actual' dropped leaf area
    # calculated dropped leaf area YY
    dropped_leaf_area(potential_leaf_area, green_leaf_area) => (potential_leaf_area - green_leaf_area) ~ track(u"cm^2")
end


"""
    Sampler(y::Vector{Int}, X::Matrix{Float64}; kwargs...)

Initialize a Gibbs Sampler for the following Bayesian Negative-Binomial model:

```math
\\begin{aligned}
y_i | x_i, \\beta, \\gamma
&\\sim
\\text{NB}(r_{0y}, 1 / (1 + e^{- x_{i\\gamma}'\\beta_{\\gamma}})),
\\\\
\\beta_k | \\gamma_k
&\\sim
\\begin{cases}
    \\mathcal{N}(\\mu_{0\\beta}, \\Sigma_{0\\beta}),
    &\\text{if } \\gamma = 1,
    \\\\
    \\delta_0
    &\\text{if } \\gamma = 0,
\\end{cases}
\\\\
\\gamma
&\\sim
\\text{Womack}(\\eta_{0\\gamma}),
\\\\
s
&\\sim
\\text{Gamma}(a_{0s}, b_{0s}),
\\end{aligned}
```

given a response vector `y` and a design matrix `X`. This constructor does not
copy any of its arguments. Hence, for example, if `y[1]` is changed, it will
affect the Gibbs sampler.

where ``y_i \\in \\mathbb{N}_0`` is the response for the ``i``-observation,
``x_i \\in \\mathbb{R}^K`` is the associated covariate vector,
``\\delta_0(\\cdot)`` is the Dirac measure centred at ``0`` and
``\\text{Womack}(\\eta_0)`` is the Womack distribution with
parameter ``\\eta_{0\\gamma}`` on ``\\{0, 1\\}^K``.

# Keyword arguments

* `β = zeros(size(X, 2))`: current state of ``\\beta``.
* `Σ0β = 10 * I(size(X, 2))`: ``\\Sigma_{0\\beta}``.
* `μ0β = zeros(size(X, 2))`: ``\\mu_{0\\beta}``.
* `a0s = 1.0`: ``a_{0s}``.
* `b0s = 1.0`: ``a_{0s}``.

# Example

```julia
julia> using Random, Distributions
julia> rng = MersenneTwister(1)
julia> X = randn(rng, 100, 2)
julia> y = rand(rng, 0:2, 10)
julia> s = BayesNegativeBinomial.Sampler(y, X)
```
"""
struct Sampler
    y::Vector{Int}
    X::Matrix{Float64}
    mapping::Vector{Vector{Int}}
    β::Vector{Float64}
    ω::Vector{Float64}
    ξ::Vector{Float64}
    ϕ::Vector{Float64}
    q::Vector{Float64}
    g::Vector{Bool}
    A::Matrix{Float64}
    b::Vector{Float64}
    s::Vector{Int}
    a0s::Float64
    b0s::Float64
    μ0β::Vector{Float64}
    Σ0β::Matrix{Float64}
    ζ0g::Float64
    update_g::Vector{Bool}
    function Sampler(
        y::Vector{Int},
        X::Matrix{Float64};
        β::Vector{Float64} = zeros(size(X, 2)),
        s::Vector{Int} = [2],
        μ0β::Vector{Float64} = zeros(size(X, 2)),
        Σ0β::Matrix{Float64} = Matrix{Float64}(10 * I(size(X, 2))),
        mapping::Vector{Vector{Int}} = [[i] for i in 1:size(X, 2)],
        a0s::Float64 = 1.0,
        b0s::Float64 = 1.0,
        ζ0g::Float64 = 1.0,
        update_g = ones(Bool, length(mapping))
    )
        N, D = size(X)
        ω = zeros(N)
        ξ = zeros(N)
        ϕ = zeros(N)
        q = zeros(N)
        g = ones(Bool, length(mapping))
        A = zeros(D, D)
        b = zeros(D)
        new(y, X, mapping, β, ω, ξ, ϕ, q, g, A, b, s, a0s, b0s, μ0β, Σ0β, ζ0g, update_g)
    end
end

"""
    sample(rng::AbstractRNG, s::BayesNegativeBinomial.Sampler; kwargs...)

Draw a posterior sample using the Gibbs sampler `s`,
following Polson et al. (2013).

# Arguments

* `mcmcsize = 2000`: the total number of MCMC iterations.
* `burnin = 1000`: the number of *burn-in* iterations.

# Example

```julia
julia> using Random
julia> rng = MersenneTwister(1)
julia> X = randn(rng, 100, 2)
julia> y = rand(rng, 0:2, 10)
julia> sampler = BayesNegativeBinomial.Sampler(y, X)
julia> chain = BayesNegativeBinomial.sample(rng, sampler)
```

# References

1. Polson, N., Scott, J. & Windle, J. (2013) Bayesian inference for logistic
    models using Pólya–Gamma latent variables, *Journal of the American
    Statistical Association*, 108:504, 1339-1349,
    <https://doi.org/10.1080/01621459.2013.829001>.
"""
function sample(rng::AbstractRNG, sampler::Sampler; mcmcsize = 10000, burnin = 5000)
    chain = [zeros(size(sampler.X, 2)) for _ in 1:(mcmcsize - burnin)]
    for iter in 1:mcmcsize
        step!(rng, sampler)
        if iter > burnin
            chain[iter - burnin] .= sampler.β
        end
    end
    return chain
end

"""
    step!(rng::AbstractRNG, s::BayesNegativeBinomial.Sampler)

Perform 1 iteration of the Gibbs sampler `s`, following Polson et al. (2013).

# Example

```julia
julia> using Random, Distributions
julia> rng = MersenneTwister(1)
julia> X = randn(rng, 100, 2)
julia> y = rand(rng, 0:2, 10)
julia> s = BayesNegativeBinomial.Sampler(y, X)
julia> BayesNegativeBinomial.step!(rng, s)
```

# References

1. Polson, N., Scott, J. & Windle, J. (2013) Bayesian inference for logistic
    models using Pólya–Gamma latent variables, *Journal of the American
    Statistical Association*, 108:504, 1339-1349,
    <https://doi.org/10.1080/01621459.2013.829001>.
"""
function step!(rng::AbstractRNG, sampler::Sampler)
    step_ϕ!(sampler)
    step_ξ!(sampler)
    # step_s!(rng, sampler)
    step_ω!(rng, sampler)
    step_A!(sampler)
    step_b!(sampler)
    step_g!(rng, sampler)
    step_β!(rng, sampler)
    return nothing
end

function step_ξ!(sampler::Sampler)
    (; ξ, X, β) = sampler
    mul!(ξ, X, β)
    return nothing
end

function step_ϕ!(sampler::Sampler)
    (; ξ, ϕ) = sampler
    @. ϕ = 1.0 / (1.0 + exp(ξ))
    return nothing
end

function step_ω!(rng::AbstractRNG, sampler::Sampler)
    (; y, ω, ξ, s) = sampler
    for i in 1:length(ω)
        ω[i] = rand(rng, PolyaGammaPSWSampler(y[i] + s[], ξ[i]))
    end
    return nothing
end

function step_g!(rng::AbstractRNG, sampler::Sampler)
    (; update_g, mapping, g, μ0β, Σ0β, ζ0g) = sampler
    D = length(μ0β)
    gexp = zeros(Bool, D)
    for d in 1:length(mapping)
        gexp[mapping[d]] .= g[d]
    end
    pg = Womack(length(g), ζ0g)
    for d in 1:length(g)
        update_g[d] || continue
        logodds = 0.0
        for val in 0:1
            g[d] = val
            gexp[mapping[d]] .= val
            m1, Σ1 = posterior_hyperparameters(sampler)
            logodds += (-1)^(val + 1) * (
                logpdf(pg, g) +
                logpdf(MvNormal(μ0β[gexp], Σ0β[gexp, gexp]), zeros(sum(gexp))) -
                logpdf(MvNormal(m1, Σ1), zeros(length(m1)))
            )
        end
        g[d] = rand(rng) < exp(logodds) / (1.0 + exp(logodds))
        gexp[mapping[d]] .= g[d]
    end
    return nothing
end

function step_β!(rng::AbstractRNG, sampler::Sampler)
    (; mapping, β, g) = sampler
    D = length(β)
    gexp = zeros(Bool, D)
    for d in 1:length(mapping)
        gexp[mapping[d]] .= g[d]
    end
    β .= 0.0
    m1, Σ1 = posterior_hyperparameters(sampler)
    β[gexp] .= rand(rng, MvNormal(m1, Σ1))
    return nothing
end

function posterior_hyperparameters(sampler::Sampler)
    (; mapping, g, A, b, μ0β, Σ0β) = sampler
    D = length(μ0β)
    gexp = zeros(Bool, D)
    for d in 1:length(mapping)
        gexp[mapping[d]] .= g[d]
    end
    Σ1 = inv(cholesky(Symmetric(A[gexp, gexp])))
    m1 = Σ1 * (b[gexp] + Σ0β[gexp, gexp] \ μ0β[gexp])
    return m1, Σ1
end

function step_A!(sampler::Sampler)
    (; X, ω, Σ0β, A) = sampler
    A .= X' * Diagonal(ω) * X + inv(Σ0β)
    return nothing
end

function step_b!(sampler::Sampler)
    (; X, y, s, b) = sampler
    b .= X' * (y .- s[]) / 2
    return nothing
end

# function step_s!(rng, sampler::Sampler)
#     (; y, q, ϕ, s, a0s, b0s) = sampler
#     N = length(q)
#     s0 = s[]
#     for i in 1:N
#         q[i] = 0
#         for j in 1:y[i]
#             q[i] += rand(rng) <= s0 / (s0 + j - 1)
#         end
#     end
#     a1s = a0s + sum(q)
#     b1s = b0s
#     for i in 1:N
#         b1s - log(1 - ϕ[i])
#     end
#     ds = Gamma(a1s, 1 / b1s)
#     s[] = rand(rng, ds)
#     return nothing
# end

struct Womack <: DiscreteMultivariateDistribution
    D::Int
    ζ::Float64
    p::Vector{Float64}
    punnormalized::Vector{Float64}
    function Womack(D::Int, ζ::Float64)
        p = big.([zeros(D); 1.0])
        for d1 in (D - 1):-1:0
            for d2 in 1:(D - d1)
                p[1 + d1] += ζ * p[1 + d1 + d2] * binomial(big(d1 + d2), big(d1))
            end
        end
        p /= sum(p)
        punnormalized = copy(p)
        for d1 in 1:D
            p[d1] /= binomial(big(D), big(d1 - 1))
        end
        return new(D, ζ, p, punnormalized)
    end
end

function pdf(d::Womack, g::Vector{Bool})
    return d.p[sum(g) + 1]
end

function logpdf(d::Womack, g::Vector{Bool})
    return log(pdf(d, g))
end


using ControlSystems, ControlSystemIdentification

logspace(a, b) = exp10.(LinRange(a, b, 50))

"""
    bodeplot(G::FOTF)

Generate the bode plot of a given FOTF object.
"""
function ControlSystems.bodeplot(G::FOTF, w = exp10.(logspace(-4, 4)))
    H1 = freqresp(w.*im, G)
    result = zeros(ComplexF64, 1, 1, length(w))
    result[1, 1, :] = H1[:]
    H1 = FRD(w, result) # Frequency response data model here
    plot(H1)
end

"""
    bode(G::FOTF)

The frequency response data model of a given FOTF object.
"""
function ControlSystems.bode(G::FOTF, w = exp10.(logspace(-4, 4)))
    H1 = freqresp(w.*im, G)
    result = zeros(ComplexF64, 1, 1, length(w))
    result[1, 1, :] .= H1[:]
    H1 = FRD(H1, w) # Frequency response data model here
    return H1
end


# Gaver-Stehfest method "Post-Widder"

# 1.
# Stehfest, Harald.
# Algorithm 368: Numerical inversion of Laplace transforms [D5].
# Communications of the ACM 13.1 (1970): 47-49.

# 2.
# Kuhlman, Kristopher L.
# Review of inverse Laplace transform algorithms for Laplace-space numerical approaches.
# Numerical Algorithms 63.2 (2013): 339-355.

# 3.
# Al-Shuaibi, Abdulaziz.
# Inversion of the laplace transform via post—widder formula.
# Integral Transforms and Special Functions 11.3 (2001): 225-232.

#= Compute Post-Widder coefficients Vk for N terms. Also referred to as Salzer summation weights or Stehfest coefficients.
Only depend on the approximation order (N) and the precision. M must be even but can be precomputed and used only once.
Coefficients get very large in magnitude and oscillate rapidly. Must be careful to avoid catastrophic cancellation for lower precision.
For  N < 18 double precision is usually ok but less accurate, utilizng defualt BigFloat precision and N = 36 usually provides sufficient accuracy.
=#
function _PWcoeffs(N::Integer)
    if isodd(N)
        N += 1
        @warn "N must be even... increment N += 1"
    end
    v = zeros(BigFloat, N)
    aux = zero(eltype(v))
    for k in 1:N
        for j in fld(k + 1, 2):minimum(Int, [k, div(N, 2)])
            aux = big(j)^(div(N, 2)) * factorial(big(2 * j))
            aux /= factorial(big(div(N, 2) - j)) * factorial(big(j)) * factorial(big(j - 1))
            aux /= factorial(big(k - j)) * factorial(big(2 * j - k))
            v[k] += aux
        end
        v[k] *= (-1)^(k + div(N, 2))
    end
    return v
end
"""
    postwid(func::Function, t::AbstractFloat, v::Array{AbstractFloat,1}=_PWcoeffs(N))

Evaluate the inverse Laplace transform of `func` at the point `t` using the Gaver-Stehfest "Post-Widder" algorithm.
Vk coefficients only depend on the number of expansion terms so can be computed once. N (which must be even) defaults to 18.
In general the accuracy in double precision is poor, so arbitrary precision is used throughout.

Increasing precision should be accompanied by an increase in the number of coefficients used.
Increasing precision without increasing number of coefficients will not yield better accuracy. The inverse is generally true as well.

This method is not robust to oscillating F(t) and must be smooth.
# Example

```jldoctest
julia> setprecision(53); InverseLaplace.postwid(s -> 1/(s + 1), 2.0)
0.1353356835639731
```
"""
function postwid(func::Function, t::AbstractFloat; v = _PWcoeffs(18))
    N = length(v)
    a = zero(eltype(v))
    for k in 1:N
        bk = convert(eltype(v), k)
        a += v[k] * func(bk * log(convert(eltype(a), 2)) / t)
    end
    return a * log(convert(eltype(a), 2)) / t
end
"""
    postwid(func::Function, t::AbstractArray, v::Array{AbstractFloat,1}=_PWcoeffs(N))

Evaluate the inverse Laplace transform of `func` over an array of `t` using the Gaver-Stehfest "Post-Widder" algorithm.
Computes coefficients once and calculates f(t) across available threads.
# Example

```jldoctest
julia> setprecision(53); postwid(s -> 1/(s + 1), 2.0:4.0)
 0.1353356835639731
 0.049786688998390505
 0.01831327193414956
```
"""
function postwid(f::Function, t::AbstractArray; v = _PWcoeffs(18))
    N = length(v)
    a = zeros(eltype(v), length(t))
    Threads.@threads for ind in eachindex(t)
        a[ind] = postwid(f, t[ind], v = v)
    end
    return a
end


using mopkg
using Test

@testset "mopkg.jl" begin
	x, f = 1000.0, x -> x^4 - 24x^3 + 7x^2 - 11x + 420
	@test minimum_nr(f, x) == (-32524.475393806228, 17.812172762347867)
	x, f = 500.0, x -> x^4 - 7x^3 + 7x^2 - 77x + 77
	@test minimum_nr(f, x) == (-387.5919553775244, 5.277905093806142)
end


module TestMin
    using StreamStats
    using Distributions
    using Base.Test

    # Min of uniform draws
    for n in rand(1:1_000_000, 100)
        xs = rand(n)
        stat = StreamStats.Min()
        for x in xs
            update!(stat, x)
        end
        online_m = minimum(stat)
        online_s = state(stat)
        online_n = nobs(stat)
        batch_m = minimum(xs)
        @test online_m == batch_m
        @test online_n == n
    end
end


# functions to support transformation of general real or complex matrices to
# Jordan normal form.
#

"""
Transform upper (quasi) triangular input matrix to block-diagonal form.
  The eigenvalues of the matrix must be ordered in a way, that multiple
  eigenvalues are stored in contiguous parts of the diagonal.
  The results stored in place of the input matrix.
"""
function blockdiagonal!(A::StridedMatrix{T}, ilo::Int, ihi::Int) where T <: Number

  while ihi > ilo
    # find lower right part of diagonal with equal eigenvalues
    isep = findlastequal(A, ilo, ihi)
    if isep > ilo
      rA = ilo:isep-1
      rB = isep:ihi
      A[rA, rB] .= sylvester(view(A, rA, rA), -view(A, rB, rB), A[rA, rB])
    end
    ihi = isep - 1
  end
  A
end

"""
Find diagonal position with first eigenvalue equal to last eigenvalue
"""
function findlastequal(A::AbstractMatrix{T}, ilo::Int, ihi::Int) where T <: Complex
  k = ihi
  d = A[k,k]
  while k > ilo && A[k-1,k-1] == d
    k -= 1
  end
  k
end

function findlastequal(A::AbstractMatrix{T}, ilo::Int, ihi::Int) where T <: Real
  k = ihi
  if ihi - ilo > 1
    if A[k,k-1] == 0
      d = A[k,k]
      while k > ilo && A[k-1,k-1] == d
        k -= 1
      end
    else
      k -= 1
      d, di = ev2(A, k)
      di > 0 || error("not expecting real eigenvalue in 2x2-block")
      while k > ilo + 1
        d2, di2 = ev2(A, k-2)
        if di2 > 0 && di2 == di && d2 == d
          k -= 2
        else
          break
        end
      end
    end
  end
  k
end

function ev2(A::AbstractMatrix{T}, k::Int) where T <: Real
  dis = discriminant(A, k, k+1)
  di = sqrt(max(-dis, zero(dis)))
  d = (A[k,k] + A[k+1,k+1]) / 2
  d, di
end


#
# Functions to work with exact multiplication and solutions for R, for its Cholesky
# (lower) decomposition L, and for L'.
#

using LinearAlgebra
using SparseArrays
using ToeplitzMatrices

"""
    convolve_same(b, kernel)

Convolve the vector `b` with `kernel`, yielding a vector of the same size
as `b`, effectively padding `b` with `length(kernel)-1` initial zeros. This is
equivalent to multiplying `b` by the banded lower triangular Toeplitz matrix
whose first column begins with `kernel`.
"""
function convolve_same(b::AbstractVector, kernel::AbstractVector)
    x = kernel[1] * Array(b)
    for i=2:length(kernel)
        x[i:end] += kernel[i] * b[1:end+1-i]
    end
    x
end

"""
    deconvolve_same(b, kernel)

Reverse the effect of `convolve_same(b, kernel)`. That is, solve `Kx=b` for
`x` where `K` is the banded lower triangular Toeplitz matrix whose first row
begins with `kernel`.
"""
function deconvolve_same(b::AbstractVector, kernel::AbstractVector)
    Nk, N = length(kernel), length(b)
    lenrek = reverse(kernel[2:end])
    x = Array(b) / kernel[1]
    for i=2:min(Nk,N)
        x[i] = (x[i]-dot(kernel[i:-1:2], x[1:i-1])) / kernel[1]
    end
    for i=Nk+1:N
        x[i] = (x[i]-dot(lenrek, x[i-Nk+1:i-1])) / kernel[1]
    end
    x
end

"""
    ARMASolver(m, N)

An object containing the information needed to perform certain noise-related
operations on a vector within an ARMA model `m::ARMAModel`, including the
following (where `R` is the ARMA model noise covariance matrix and `L` is its
lower Cholesky factor):

1. L\\v  `whiten(solver, v)`
1. L*v  `unwhiten(solver, v)`
1. R\\v  `solve_covariance(solver, v)`
1. R*v  `mult_covariance(solver, v)`
1. inv(R[1:n,1:n]) `inverse_covariance(solver, n)`

Note that the object is constructed with a maximum vector length `N`, so that
these operations require `length(v) <= N` and `n <= N`."""
mutable struct ARMASolver
    p         ::Int
    q         ::Int
    phicoef   ::Vector{Float64}
    covarIV   ::Vector{Float64}
    # RR = Phi*R*Phi', a banded matrix = RRt + RRu
    RRu       ::Matrix{Float64}   # The upper left (p)x(p+q) corner of RR-RRt
    RRt       ::Vector{Float64}   # The banded Toeplitz part of RR
    LL        ::SparseMatrixCSC{Float64}   # LL = lower Cholesky factor of RR, also banded

    function ARMASolver(p, q, phicoef, covarIV, RRu, RRt, LL)
        @assert p>=0
        @assert q>=0
        @assert p+1 == length(phicoef)
        @assert length(covarIV) >= 1+q
        @assert length(covarIV) >= p
        @assert phicoef[1] == 1.0

        new(p, q, phicoef, covarIV, RRu, RRt, LL)
    end
end

function ARMASolver(m::ARMAModel, N::Integer)
    covar = model_covariance(m, 1+m.p+2m.q)
    R_corner  = SymmetricToeplitz(covar)
    Nc = length(covar)
    Nbands = max(m.p, m.q+1)
    x = fill(0.0, Nc)
    y = fill(0.0, Nc)
    x[1:m.p+1] = m.phicoef
    y[1] = x[1]
    Phi = Toeplitz(x, y)
    # Hermitian enforces exact symetry even after round-off error.
    RR_corner = Hermitian(Phi*R_corner*Phi')
    RR_toeplitz = fill(0.0, 2m.q+1)
    RR_toeplitz[1:m.q+1] = RR_corner[end, end-m.q:end]
    RR_toeplitz[m.q+1:end] = RR_corner[end, end:-1:end-m.q]
    RR_rectcorner = RR_corner[1:max(m.p,Nbands-1), :]

    # Now compute LL such that LL*LL' == RR, without ever representing RR as a full
    # matrix in memory. Awesome.
    #
    # Find the nontrivial sparse elements of this lower-triangular banded matrix
    # LL, then construct with zero values, then fill them.
    Ndata = N*Nbands - div(Nbands*(Nbands-1),2)
    rows = Array{Int}(undef, Ndata)
    cols = Array{Int}(undef, Ndata)
    vals = fill(0.0, Ndata)
    i = 1
    for r=1:min(N)
        for c=1+max(0,r-Nbands):r
            rows[i] = r
            cols[i] = c
            i += 1
        end
    end
    @assert i-1 == Ndata
    LL = sparse(rows, cols, vals, N, N)

    # Fill LL with its nontrivial values.
    LL_corner = cholesky(RR_corner).L
    for r=1:min(Nc,N)
        for c=1+max(0,r-Nbands):r
            LL[r,c] = LL_corner[r,c]
        end
    end
    for r = Nc+1:N
        mincol = r-m.q
        LL[r,mincol] = RR_toeplitz[1] / LL[mincol,mincol]
        for c=mincol+1:r-1
            LL[r,c] = (RR_toeplitz[c-r+Nbands] - dot(vec(LL[r,mincol:c-1]), vec(LL[c,mincol:c-1]))) / LL[c,c]
        end
        Rt = RR_toeplitz[m.q+1]
        S = sum( LL[ r,mincol:r-1] .^ 2)
        LL[r,r] = sqrt(Rt - S)
    end
    ARMASolver(m.p, m.q, m.phicoef, covar, RR_rectcorner, RR_toeplitz, LL)
end

"""
    whiten(solver::ARMASolver, v::Vector)

Use `solver` for an ARMA model to whiten the vector `v`. Here, "whiten" means
return `w=L\\v`, where `L` is the lower Cholesky factor of the covariance
matrix. If `v` has expectation zero, then  the expected value of `w*w'` over
noise realizations is the identity matrix.

In place of `v`, a matrix can be used. Its columns will be whitened.
"""
function whiten(solver::ARMASolver, v::AbstractVector)
    Phiv = convolve_same(v, solver.phicoef)
    nv = length(v)
    if nv  < size(solver.LL)[1]
        return solver.LL[1:nv, 1:nv] \ Phiv
    end
    solver.LL \ Phiv
end

function whiten(solver::ARMASolver, M::AbstractMatrix)
    ws(v::AbstractVector) = whiten(solver, v)
    mapslices(ws, M, dims=1)
end

"""
    unwhiten(solver::ARMASolver, w::AbstractVector)

Use `solver` for an ARMA model to unwhiten the vector `w`. Here, "unwhiten"
means return `v=L*w`, where `L` is the lower Cholesky factor of the covariance
matrix. If the expected value of `w` is zeros and of `w*w'` is the identity
matrix, then the expected value of `v*v'` is the data covariance matrix of the
ARMA model.

In place of `v`, a matrix can be used. Its columns whill be unwhitened.
"""
function unwhiten(solver::ARMASolver, w::AbstractVector)
    nw = length(w)
    if nw  < size(solver.LL)[1]
        x = solver.LL[1:nw, 1:nw] * w
    else
        x = solver.LL * w
    end
    deconvolve_same(x, solver.phicoef)
end

function unwhiten(solver::ARMASolver, M::AbstractMatrix)
    uws(v::AbstractVector) = unwhiten(solver, v)
    mapslices(uws, M, dims=1)
end

"""
    mult_covariance(solver::ARMASolver, v::AbstractVector)

Use `solver` for an ARMA model to multiply the vector `v` by the covariance matrix.
"""
function mult_covariance(solver::ARMASolver, v::AbstractVector)
    # Reversing before and after deconvolve ensures that we are solving Phi', not Phi.
    v1 = reverse(deconvolve_same(reverse(v), solver.phicoef))

    # Multiply by RR, the "moving-average-only R matrix".
    Nv = length(v1)
    Nu1,Nu2 = size(solver.RRu)
    q = solver.q
    if Nv > Nu2
        x = similar(v1)
        x[1:Nu1] = solver.RRu * v1[1:Nu2]
        for r = Nu1+1 : Nv-q
            x[r] = dot(solver.RRt, v1[r-q:r+q])
        end
        for r = Nv-q+1 : Nv
            x[r] = dot(solver.RRt[1:end-r+Nv-q], v1[r-q:end])
        end

    elseif Nv > Nu1
        x = similar(v1)
        x[1:Nu1] = solver.RRu[:, 1:Nv] * v1
        for r = Nu1+1 : Nv-q
            x[r] = dot(solver.RRt, v1[r-q:r+q])
        end
        for r = max(Nu1,Nv-q)+1 : Nv
            x[r] = dot(solver.RRt[1:end-r+Nv-q], v1[r-q:end])
        end

    else
        x = solver.RRu[1:Nv,1:Nv] * v1
    end
    deconvolve_same(x, solver.phicoef)
end

function mult_covariance(solver::ARMASolver, M::AbstractMatrix)
    mc(v::AbstractVector) = mult_covariance(solver, v)
    mapslices(mc, M, dims=1)
end

"""
    solve_covariance(solver::ARMASolver, v::AbstractVector)

Use `solver` for an ARMA model to solve `Rx=v` where `R` is the noise covariance
matrix.
"""
function solve_covariance(solver::ARMASolver, v::AbstractVector)
    nv = length(v)
    v1 = whiten(solver, v)
    if nv  < size(solver.LL)[1]
        v2 = solver.LL[1:nv,1:nv]' \ v1
    else
        v2 = solver.LL' \ v1
    end
    reverse(convolve_same(reverse(v2), solver.phicoef))
end

function solve_covariance(solver::ARMASolver, M::AbstractMatrix)
    sc(v::AbstractVector) = solve_covariance(solver, v)
    mapslices(sc, M, dims=1)
end

"""
    inverse_covariance(solver::ARMASolver, N)

Use `solver` for an ARMA model to compute the inverse noise covariance matrix,
to size `[N,N]`.
"""
function inverse_covariance(solver::ARMASolver, N::Integer)
    M = Matrix(1.0I, N, N)
    hcat([solve_covariance(solver, M[:,i]) for i=1:N]...)
end


g = random_regular_graph(20, 3)
J = 100
h = rand(20)
σ, E = ground_state_mincut(g, h, J)
@test σ == ones(20)
σ, E = ground_state_mincut(g, -h, J)
@test σ == -ones(20)

J = Vector{Vector{Float64}}()
for i=1:nv(g)
    push!(J, zeros(degree(g, i)))
    for (k, j) in enumerate(neighbors(g, i))
        J[i][k] = rand()
    end
end

h = randn(20)
σ, E = ground_state_mincut(g, h, J)
σ2, E2 = ground_state_mincut(g, 2 + h, J)
@test all(σ .<= σ2)


@testset "MST" begin
    using STMO
    using STMO.MST

    edges = [(1, :A, :D),
            (2, :A, :B),
            (3, :D, :C),
            (4, :B, :D),
            (5, :B, :C)]

   vertices = [:A, :B, :C, :D]

   mst_prim, cprim = prim(vertices, edges, :A)
   @test cprim == 6
   @test (4, :B, :D) ∉ mst_prim && (4, :D, :B) ∉ mst_prim

   @test prim(vertices, edges)[2] == 6

   mst_kruskal, ckruskal = kruskal(vertices, edges)
   @test ckruskal == 6
   @test (4, :B, :D) ∉ mst_kruskal && (4, :D, :B) ∉ mst_kruskal
end


using LinearAlgebra
using Statistics
using StaticArrays

@testset "Geometry" begin

for type_under_test in [Vector3, Point3, Normal3]
    typename = string(type_under_test)
    @testset "Common tests: $typename" begin
        X = type_under_test(1.0, 0.0, 0.0)
        Y = type_under_test(0.0, 1.0, 0.0)
        Z = type_under_test(0.0, 0.0, 1.0)
        @testset "Creating" begin
            v = type_under_test(1.0, 2.0, 3.0)
            @test v.x ≈ 1.0
            @test v.y ≈ 2.0
            @test v.z ≈ 3.0
            v = type_under_test(1,2,3)
            @test v.x ≈ 1.0
            @test v.y ≈ 2.0
            @test v.z ≈ 3.0
            v = type_under_test(1//2, 3//4, -2//1)
            @test v.x ≈ 0.5
            @test v.y ≈ 0.75
            @test v.z ≈ -2.0
            v = type_under_test(1.0)
            @test v.x ≈ 1.0
            @test v.y ≈ 1.0
            @test v.x ≈ 1.0
            v = type_under_test(Vector3(1, 2, 3))
            @test v.x ≈ 1.0
            @test v.y ≈ 2.0
            @test v.z ≈ 3.0
        end

        @testset "Arithmetic" begin
            v = type_under_test(1, 2, 3)
            @test 2.0 * v ≈ type_under_test(2, 4, 6)
            @test v * 2.0 ≈ type_under_test(2, 4, 6)
            @test v/2 ≈ type_under_test(0.5, 1.0, 1.5)
        end

        @testset "Normalizing" begin
            v = type_under_test(2, 0, 0)
            @test normalize(v) ≈ type_under_test(1, 0, 0)
            v = type_under_test(0, -2, 0)
            @test normalize(v) ≈ type_under_test(0, -1, 0)
            v = type_under_test(3, -4, 0)
            @test normalize(v) ≈ type_under_test(3//5, -4//5, 0)
        end

        @testset "Dot product" begin
            v = type_under_test(0.5, 1.0, -2.0)
            @test dot(v, X) ≈ 0.5
            @test dot(v, Y) ≈ 1.0
            @test dot(v, Z) ≈ -2.0

            u = type_under_test(2, 1, 1)
            @test dot(v, u) ≈ 0.0
        end

        @testset "Min/max" begin
            v = type_under_test(-1, 0, 2)
            u = type_under_test(0, -1, 0)
            @test min(v,u) ≈ type_under_test(-1, -1, 0)
            @test max(u,v) ≈ type_under_test(0, 0, 2)
        end

        @testset "Iteration" begin
            V = type_under_test(1, 2, 3)
            A = [x for x in V]
            @test A ≈ [1.0, 2.0, 3.0]
        end

        @testset "Cross product" begin
            X = type_under_test(1,0,0)
            Y = type_under_test(0,1,0)
            Z = type_under_test(0,0,1)
            @test cross(X, Y) ≈ Vector3(0, 0, 1)
            @test cross(X, Z) ≈ Vector3(0, -1, 0)
            @test cross(Y, X) ≈ Vector3(0, 0, -1)
            @test cross(Y, Z) ≈ Vector3(1, 0, 0)
            @test cross(Z, X) ≈ Vector3(0, 1, 0)
            @test cross(Z, Y) ≈ Vector3(-1, 0, 0)
        end

        @testset "Various functions" begin
            X = type_under_test(1, 2, 3)
            @test size(X) == 3
            @test length(X) == 3
            @test !isnan(X)
            @test isnan(type_under_test(NaN, 2, 3))
            @test isnan(type_under_test(1, NaN, 3))
            @test isnan(type_under_test(1, 2, NaN))
        end

        @testset "Conversions" begin
            T = type_under_test(1, 2, 3)
            @test convert(Array{Float64, 1}, T) ≈ [1.0, 2.0, 3.0]

            V = Vector3(1, 2, 3)
            P = Point3(1, 4, 9)
            N = Normal3(1, 8, 27)
            @test convert(type_under_test, V) ≈ type_under_test(1, 2, 3)
            @test convert(type_under_test, P) ≈ type_under_test(1, 4, 9)
            @test convert(type_under_test, N) ≈ type_under_test(1, 8, 27)
        end
    end
end

@testset "Type-specific functions" begin
    P1 = Point3(2, 0, -4)
    P2 = Point3(0, 0, -2)
    @test mean(P1, P2) ≈ Point3(1, 0, -3)

    V1 = Vector3(1, 2, 3)
    V2 = Vector3(1, 0, -1)
    @test V1 + V2 ≈ Vector3(2, 2, 2)
    @test V1 - V2 ≈ Vector3(0, 2, 4)
end

@testset "Transformations" begin
    @testset "Identity transformation" begin
        T = Transformation()
        @test T.M ≈ Matrix{Float64}(I, 4, 4)
        @test T.MInv ≈ Matrix{Float64}(I, 4, 4)
    end
    @testset "Translation" begin
        point = Point3(0)
        T = translation(1, 0, 0)
        @test T(point) ≈ Point3(1.0, 0.0, 0.0)
        U = inv(T)
        @test U(T(point)) ≈ point

        point = Point3(-1.0, 0.0, 2.0)
        T = translation(1, 2, 5)
        @test T(point) ≈ Point3(0.0, 2.0, 7.0)
        U = inv(T)
        @test U(T(point)) ≈ point
    end
    @testset "Rotation" begin
        X = Vector3(1, 0, 0)
        Y = Vector3(0, 1, 0)
        Z = Vector3(0, 0, 1)

        point = Point3(1,0,0)
        T = rotation(X, pi/2)
        @test T(point) ≈ Point3(1.0, 0.0, 0.0)
        U = inv(T)
        @test U(T(point)) ≈ point

        point = Point3(1.0, 0.0, 0.0)
        T = rotation(Z, pi/2)
        @test isapprox(T(point), Point3(0.0, 1.0, 0.0) ; atol=1e-16)
        U = inv(T)
        @test U(T(point)) ≈ point
        T = rotation(Y, pi/2)
        @test isapprox(T(point), Point3(0.0, 0.0, -1.0) ; atol=1e-16)
        U = inv(T)
        @test U(T(point)) ≈ point
    end
    @testset "Scaling" begin
        P = Point3(-1, 1, 2)
        T = scaling(2, 1, 1)
        @test T(P) ≈ Point3(-2, 1, 2)
        U = inv(T)
        @test U(T(P)) ≈ P
        T = scaling(1, 2, 1)
        @test T(P) ≈ Point3(-1, 2, 2)
        U = inv(T)
        @test U(T(P)) ≈ P
        T = scaling(1, 1, 2)
        @test T(P) ≈ Point3(-1, 1, 4)
        U = inv(T)
        @test U(T(P)) ≈ P

    end
    @testset "Miscellaneous" begin
        v = Vector3(1/sqrt(3), 1/sqrt(3), 1/sqrt(3))
        for direction in [X_AXIS, Y_AXIS, Z_AXIS, v]
            T = GRAYTR.rotate_z_to(direction)
            @test isapprox(T(Z_AXIS), direction ; atol=1e-16)
        end
    end
end

end # testset Geometry


@testset "DFS" begin
    pred1 = Predicate((x,y) -> x == -y || x == y)
    indl1 = -2:3
    pred1_adj = build_adj_matrix(pred1, indl1)
    eqC_1 = EquivalenceClasses(pred1, indl1)
end

@testset "Adjacency Matrix" begin
    pred1 = Predicate((x,y) -> x == -y)
    pred2 = Predicate((x,y)->x==y+3)
    indl1 = -2:3
    indl2 = -3:5
    eqC_1 = EquivalenceClasses(pred1, indl1)
    adj2 = build_adj_matrix(pred2, indl2)
    @test all(build_adj_matrix(pred1, indl1) .== Bool[1 0 0 0 1 0; 0 1 0 1 0 0; 0 0 1 0 0 0; 0 1 0 1 0 0; 1 0 0 0 1 0; 0 0 0 0 0 1])
    #@test all(find_classes(adj2) == find_classes(pred2, indl2))
end

@testset "Find Classes Adj" begin
    adj1 = falses(5,5)
    adj2 = trues(5,5)
    adj3 = convert(BitArray, [0 1 0; 1 0 0; 0 0 0])
    @test all(find_classes(adj1) .== collect(1:5))
    @test all(find_classes(adj2) .== ones(5))
    @test all(find_classes(adj3) .== UInt32.([1, 1, 3]))
end

@testset "Find Classes Paths" begin
    m1(x) = [-x, x+3]
    cl, ops = find_classes(m1,[0x00001, 0x00002],-3:5)
    println("---")
    println(cl)
    println(ops)
    println("---")
end


## load packages
using Plots, LaTeXStrings, VortexDistributions
gr(titlefontsize=12,size=(500,300),transpose=true,colorbar=false)

function psimovie(sol,sim)
    @unpack X,t = sim; x,y = X
    anim = @animate for i in eachindex(t)
        ψ = xspace(sol[i],sim)
        showpsi(x,y,ψ)
    end
    return anim
end

## Set up parameters and ground state
# Units of ξ for length, 1/μ for time
# related by μ ≡ ħ²/mξ² = mc^2 where
# c = ħ/mξ is the speed of sound of
# the uniform system.

# Initialize simulation
L = (200.,200.)
N = (512,512)
sim = Sim(L,N)
@unpack_Sim sim

# set simulation parameters
μ = 1.0
g = 0.01
γ = 0.5
ti = 0.0
tf = .3pi
Nt = 150
t = LinRange(ti,tf,Nt)

# Useful state functions
ψ0(x,y,μ,g) = sqrt(μ/g)*sqrt(max(1.0-V(x,y,0.0)/μ,0.0)+im*0.0)
healinglength(x,y,μ,g) = 1/sqrt(g*abs2(ψ0(x,y,μ,g)))

## Make initial state
x,y = X
kx,ky = K
ψi = ψ0.(x,y',μ,g)
ϕi = kspace(ψi,sim)
@pack_Sim! sim

## evolve
sol = runsim(sim)

## plot ground state
ϕg = sol[end]
ψg = xspace(ϕg,sim)
showpsi(x,y,ψg)

## periodic dipole phase
# Billam et al, PRL 112, 145301 (2014), Supplemental
# Methods
H(x) = x > 0. ? 1.0 : 0.0
shift(x,xi) = x - xi
tans(x,xk) = tan((shift(x,xk) - π)*0.5)
tanhs(x,xk,j) = tanh((shift(x,xk) + 2*π*j)*0.5)

function kernel(x,y,xp,yp,xn,yn,j)
    return atan(tanhs(y,yn,j)*tans(x,xn)) -
    atan(tanhs(y,yp,j)*tans(x,xp))
end

# Dimensionless form
function θd(x,y,dip)
    vp,vn = dip
    xp,yp,_ = rawData(vp)
    xn,yn,_ = rawData(vn)
    s = 0.0
    for j = -5:5
        s += kernel(x,y,xp,yp,xn,yn,j)
    end
    return s + π*(H(shift(x,xp)) - H(shift(x,xn))) - y*(xp - xn)/(2*π)
end
# arbitrary domains and dipole sizes:
function thetad(x,y,xp,yp,xn,yn)
    s = 0.0
    for j = -5:5
        s += kernel(x,y,xp,yp,xn,yn,j)
    end
    s += π*(H(shift(x,xp)) - H(shift(x,xn))) - y*(xp - xn)/(2*π)
    return s - x*H(abs(yp - yn) - π) + y*H(abs(xp - xn) - π)
end

function Thetad(x,y,xp,yp,xn,yn)
    Lx = x[end]-x[1]
    Ly = y[end]-y[1]
    return @. angle(exp(im*thetad.(x*2*pi/Lx,y'*2*pi/Ly,xp*2*pi/Lx,yp*2*pi/Ly,xn*2*pi/Lx,yn*2*pi/Ly)))
end

## initial dipole, with periodic phase
ψv = copy(ψg)
d = 14
ξv = healinglength(0.,0.,μ,g)
# can use default ξ = 1
xp,yp = -50.,d/2
xn,yn = -50.,-d/2
pv = PointVortex(xp,yp,1)
nv = PointVortex(xn,yn,-1)

dipole = [nv;pv]
psi = Torus(ψv,x,y) # set field topology for VortexDistributions
vortex!(psi,dipole) # make dipole
ψv = abs.(psi.ψ).*exp.(im*Thetad(x,y,xp,yp,xn,yn))
showpsi(x,y,ψv)

## make sure phase is periodic
using Test
testphase = angle.(ψv)
@test testphase[:,1] ≈ testphase[:,end]
@test testphase[1,:] ≈ testphase[end,:]

## evolve dipole with weak damping
# Set simulation parameters
c = sqrt(μ)
tf = 3*L[1]/c # a short run for testing evolutoin and detection
γ = 0.03

t = LinRange(ti,tf,Nt)
ϕi = kspace(ψv,sim)
alg = Vern7()

@pack_Sim! sim

# remove boundary artifacts with small γ
@time solv = runsim(sim)

## plot and animate
ψd = xspace(solv[end],sim)
showpsi(x,y,ψd)

anim = psimovie(solv,sim)
gif(anim,"./examples/dipole_damping.gif",fps=25)

## Vortex detection test
Nt = 100
d = zeros(Nt)

for i=1:Nt
    psi = Torus(xspace(solv[i],sim),x,y)
    vort = findvortices(psi) |> rawData
    d[i] = abs(vort[2,2]-vort[1,2])
end

plot(t[1:Nt],d)

## Hamiltonian evolution
γ = 0.0
tf = L[1]/c/5
t = LinRange(ti,tf,Nt)
ϕi = kspace(ψv,sim)
@pack_Sim! sim

solh = runsim(sim)

## plot
ψdh = xspace(solh[end],sim)
showpsi(x,y,ψdh)

anim = psimovie(solh,sim)
gif(anim,"./examples/dipole_hamiltonian.gif",fps=25)

## energy densities
function showenergies(ψ)
    et,ei,ec = energydecomp(ψ)
    p1 = heatmap(x,y,log10.(ei),aspectratio=1)
    xlabel!(L"x/\xi");ylabel!(L"y/\xi")
    title!("Incompressible")
    p2 = heatmap(x,y,log10.(ec),aspectratio=1)
    xlabel!(L"x/\xi");ylabel!(L"y/\xi")
    title!("Compressible")
    p = plot(p1,p2,size=(600,300))
    return p
end

K2 = k2(K)

anim = @animate for i in eachindex(t)
    ψ = xspace(solh[i],sim)
    psi = XField(ψ,X,K,K2)
    showenergies(psi)
end

gif(anim,"./examples/dipoleenergies.gif",fps=30)

## energy totals

function venergy(ϕ,sim,t)
    @unpack g,X = sim; x,y = X
    ψ = xspace(ϕ,sim)
    @. ψ *= V(x,y',t)
    return kspace(ψ,sim)
end

function ienergy(ϕ,sim)
    @unpack g,X = sim; x,y = X
    ψ = xspace(ϕ,sim)
    @. ψ *= 0.5*g*abs2(ψ)
    return kspace(ψ,sim)
end

function kinetic(ϕ,sim)
    @unpack espec,K = sim; kx,ky = K
    dkx,dky = kx[2]-kx[1],ky[2]-ky[1]
    return sum(espec.*abs2.(ϕ))*dkx*dky |> real
end

dx = diff(x)[1]; dy = diff(y)[1]
dkx = diff(kx)[1]; dky = diff(ky)[1]
n0 = μ/g

function energies(sol)
    t = sol.t
    # H = zero(t)
    Ekall = zero(t)
    Ekhy = zero(t)
    Ei = zero(t)
    Ec = zero(t)
    Ev = zero(t)
    Eint = zero(t)
    Eqp = zero(t)
    Et = zero(t)
    Natoms = zero(t)

    for (i,t) in enumerate(t)
    ϕ = sol[i]
    ψ = xspace(ϕ,sim)
    psi = XField(ψ,X,K,K2)
    ek,ei,ec = energydecomp(psi)
    Ekhy[i] = sum(ek)*dx*dy |> real
    Ei[i] = sum(ei)*dx*dy |> real
    Ec[i] = sum(ec)*dx*dy |> real
    Ekall[i] = kinetic(ϕ,sim)

    Ev[i] = sum(@. V(x,y',t)*abs2(ψ))*dx*dy |> real
    Eint[i] = sum(@. g/2*(abs2(ψ) - n0)^2)*dx*dy |> real
    Eqp[i] = Ekall[i] - Ekhy[i]

    Et[i] = Ekall[i] + Ev[i] + Eint[i]
    Natoms[i] = sum(abs2.(ψ))*dx*dy
    end
    return Ei,Ec,Ekhy,Eqp,Ev,Eint,Et,Ekall,Natoms
end

@time Ei,Ec,Ekhy,Eqp,Ev,Eint,Et,Ekall,Natoms = energies(solh)

## plot energies
# simple breakdown
plot(t,Et./Natoms,label=L"E_t",legend=:left)
plot!(t,Eint./Natoms,label=L"E_{int}")
plot!(t,Ekall./Natoms,label=L"E_{kall}")
ylims!(0,1.1*Et[1]./Natoms[1])
xlabel!(L"t")

# decomposition
plot(t,(Ekall .+ Eint)./Natoms,label=L"E_{kall}+E_{int}",w=5,c=:gray,alpha=0.3)
plot!(t,Et./Natoms,label=L"E_t",legend=:left)
plot!(t,Ei./Natoms,label=L"E_k^i")
plot!(t,Ec./Natoms,label=L"E_k^c")
plot!(t,Eint./Natoms,label=L"E_{int}")
plot!(t,Eqp./Natoms,label=L"E_{qp}")
plot!(t,Ekall./Natoms,label=L"E_{kall}")
xlabel!(L"t")

# if we need to save data:
# @save "./examples/dipoledecay.jld2" solv.u sim

# another example for saving data as individual files
# using a call back in OrdinaryDiffEq


using FresnelIntegrals
using Test
using QuadGK

@testset "FresnelIntegrals.jl" begin
    # Generate random complex number
    z = randn(ComplexF64)
    # Test by comparing to numeric solution
    @test fresnelc(z) ≈ quadgk(t->cos(π*t^2/2),0,z)[1]
    @test fresnels(z) ≈ quadgk(t->sin(π*t^2/2),0,z)[1]
    # Test just for code coverage 😄
    @test (fresnelc(z),fresnels(z)) == fresnel(z)
end


using Distributions
using Calculus
using OrdinaryDiffEq
using Plots
gr()

# Force constant from Wang = 14500 amu/ps^2 = 14500 * 1.66e-27(kg) / 10^(-24)(s^2) = 24.07
# equilibrium bond length for hydrogen = 7.5e-11m
# energy = 4.5 eV = 7.2e-19
# coupling constant = 3500 cm−1/Å = 3500 *100*10^10 = 3500 e-8 = 3.5 e-5

#
# ---------------
# constants used 04/08/18:
# J_if = 0.136056925
# en = 122.4512325 (but plotted with en = 0)
# R_0 = 1.4423076923076923
# K = 1.3204114285714283
# ----------------

const M = 1

"""
Function to define the non-adiabatic potential for a single site centered on R_0
Basic quadratic approximation is used with minimum energy e_0
"""
function non_adiabatic_potential(K, R_0, e_i0, e_f0)
    n_a_p_i = function (R) return (K*(R+R_0/2)^2 + e_i0) end
    n_a_p_f = function (R) return (K*(R-R_0/2)^2 + e_f0) end
    return n_a_p_i, n_a_p_f
end

"""
Function to define the adiabatic potentials from the non_adiabatic_potentials
"""
function adiabatic_potential(K, R_0, e_i0, e_f0, J_if)
    n_a_p_i, n_a_p_f = non_adiabatic_potential(K, R_0, e_i0, e_f0)
    a_p_m = function (R) return 0.5*(n_a_p_i(R)+n_a_p_f(R)) - 0.5*sqrt((n_a_p_i(R)-n_a_p_f(R))^2+4*J_if^2) end
    a_p_p = function (R) return 0.5*(n_a_p_i(R)+n_a_p_f(R)) + 0.5*sqrt((n_a_p_i(R)-n_a_p_f(R))^2+4*J_if^2) end
    return a_p_m, a_p_p
end

"""
Function to plot diabatic and adiabatic potentials for simple 2-atom system.
"""
function plot_potentials(K, R_0, e_i0, e_f0, J_if)
    plot(r -> non_adiabatic_potential(K, r, R_0, e_i0, e_f0)[1], -R_0, R_0, color = :red, label = "e_i")
    plot!(r -> non_adiabatic_potential(K, r, R_0, e_i0, e_f0)[2], -R_0, R_0, color = :red, label = "e_f")
    plot!(r -> adiabatic_potential(K, r,R_0, e_i0, e_f0, J_if)[1], -R_0, R_0, color = :blue, label = "e_+")
    plot!(r -> adiabatic_potential(K, r,R_0, e_i0, e_f0, J_if)[2], -R_0, R_0, color = :blue, label = "e_-")
    gui()
end

"""
Function to calculate the force field from the potential energy curves, V.
using the relation F = m*a = -∇V
"""
function force_from_V(a_p_m, a_p_p, plotting::Bool=true)
    F_p(x) = -derivative(R -> a_p_p(R),x)
    F_m(x) = -derivative(R -> a_p_m(R),x)

    if plotting
        plot(x->F_p(x), -R_0, R_0, label = "F_p")
        plot!(x->F_m(x), -R_0, R_0, label = "F_m")
        gui()
    end
    return F_m, F_p
end

"""
Function to calculate the second derivative of the potential (i.e the slope of the
force) at every point in space to be used to determine a classical propagation.
"""
function slope(F_m)
    slope_m(x) = derivative(R -> F_m(R),x)
end

"""
Function to define the equation of motion of adiabatic surface
"""
function Acceleration(du,u,p,t,F)
   x = u[1]
   dx = u[2]
   du[1] = dx
   du[2] = real(F(x/2)/M)
end

"""
choice of T and dt are important.
"""
function classical_propagation_2(x_0, v_0,T::Float64, dt, F)
    #F = function (R) return F_m(R/2) end
    tspan = (0.0,T)
    myAcceleration(du,u,p,t)=Acceleration(du,u,p,t,F)

    EoM = ODEProblem(myAcceleration,[x_0,v_0],tspan)
    sol = solve(EoM, Euler(), dt = dt)
    x = sol[1,:]
    v = sol[2,:]
    return x,v
end

"""
Function to trace motion of atom
x_0 = initial displacement from equilibrium bond length.
"""
function plot_trajectory(dx::Float64, v_0::Float64, K, R_0, e_i0, e_f0, m, dt,T::Float64)
    x = zeros(T/dt); v = zeros(T/dt); V1 = zeros(T/dt); V2 = zeros(T/dt); t = zeros(T/dt)
    x_f0 = (R_0+dx)/2; x_i0 = -x_f0; x_eq = R_0;
    x[1] = R_0 + dx; v[1] = v_0; t[1] = dt
    n_a_p_i = non_adiabatic_potential(K, R_0, e_i0, e_f0)[1]
    n_a_p_f = non_adiabatic_potential(K, R_0, e_i0, e_f0)[2]

    for i in 1:Int(T/dt-1)
        phis(R) = non_adiabatic_states(R, [-x[i]/2,x[i]/2])
        phi_1(R)= phis(R)[1]; phi_2(R) = phis(R)[2]
        J_if = overlap(phi_1,phi_2)
        a_p_m = adiabatic_potential(K, R_0, e_i0, e_f0, J_if)[1]
        a_p_p = adiabatic_potential(K, R_0, e_i0, e_f0, J_if)[2]

        F_m, F_p = force_from_V(a_p_m, a_p_p, false)

        x_new, v_new = classical_propagation_2(x[i], v[i], dt, dt, F_m)
        x[i+1],v[i+1] = x_new[end], v_new[end]
        V1 = a_p_m(-x[i+1]/2); V2 = a_p_p(-x[i+1]/2); t[i+1] = (i+1)*dt

        if i%10 == 0
            p1 = plot(n_a_p_i,-R_0,R_0, color = :blue)
            p1 = plot!(n_a_p_f,-R_0,R_0,color = :blue)
            p1 = plot!(a_p_m,-R_0,R_0,color = :red)
            p1 = plot!(a_p_p,-R_0,R_0,color = :red)
            p1 = scatter!([-x[i+1]/2], [a_p_m(-x[i+1]/2)],color = :black, markershape = :circle)
            p1 = scatter!([x[i+1]/2], [a_p_m(x[i+1]/2)],color = :red, markershape = :circle)
            gui()
        end

    end
    x1 = -x/2; x2 = x/2;
    KE = 0.5*v.^2
    PE = 2*V1
    TE = KE + PE
    # for i in 1:length(x)
    #     if i%100==0
    #         p1 = plot(r->a_p_m(r),-R_0,R_0)
    #         p1 = scatter!([x1[Int(i)]],[V1[Int(i)]],color = :black, markershape = :circle)
    #         gui()
    #     end
    # end

    p2 = plot(x, KE)
    p2 = plot!(x, PE)
    p2 = plot!(x, TE)
    p3 = plot(t, x)
    plot(p2,p3,layout = (1,2))
    gui()
    #println("min x = ",xmin, "F = ", f_xmin, "max x = ", xmax, "F = ", f_xmax )
    return x, v
end

"""
Function to calculate eigenstates
create Hamiltonian from the onsite energies (e_i and e_f) calculated and
coupling constants J.
find adiabatic states as eigenvectors of this Hamiltonian.
"""
# function adiabatic_states(x, a_p_m, a_p_p)
#     E_11 = a_p_m(x)
#     E_22 = a_p_p(x)
#     H = [E_11 0; 0 E_22]
#     psis = eigvecs(H)
#     return psis
# end

"""
Generate non adiabatic potentials as functions of R.
Find eigenvectors (adiabatic states) from the eigenstates.
Plot eigenstate as a function of nuclear coordinates.
    (Figures show excited state )
"""
function plot_states(K, R_0, dR, e_i0, e_f0, J_if)
    plot()
    x = R_0+dR
    R_n1 = R_e = x/2
    R_n2 = -x/2
    n_a_p_i = non_adiabatic_potential(K, R_0, e_i0, e_f0)[1]
    n_a_p_f = non_adiabatic_potential(K, R_0, e_i0, e_f0)[2]
    a_p_p = adiabatic_potential(K, R_0, e_i0, e_f0, J_if)[1]
    a_p_m = adiabatic_potential(K, R_0, e_i0, e_f0, J_if)[2]
    H_diabatic = [n_a_p_i(R_e) J_if; J_if n_a_p_f(R_e)]
    H2_diabatic = [n_a_p_i(-R_e) J_if; J_if n_a_p_f(-R_e)]
    #H_adiabatic = [a_p_m(R_e) 0; 0 a_p_p(R_e)p3 = ]
    cr_11, cr_12, cr_21, cr_22 = eigvecs(H_diabatic)
    cl_11, cl_12, cl_21, cl_22 = eigvecs(H2_diabatic)
    phis(R) = non_adiabatic_states(R, [R_n1, R_n2])
    psir_1(R) = cr_11*phis(R)[1] + cr_12*phis(R)[2]
    psir_2(R) = cr_21*phis(R)[1] + cr_22*phis(R)[2]
    psil_1(R) = cl_11*phis(R)[1] + cl_12*phis(R)[2]
    psil_2(R) = cl_21*phis(R)[1] + cl_22*phis(R)[2]
    # Psis(R) = adiabatic_states(R, a_p_m, a_p_p)
    # Phis(R) = non_adiabatic_states(R, [-x/2,x/2])
    # psi_1(R) = Psis(R)[:,1][1]*Phis(R)[1] + Psis(R)[:,1][2]*Phis(R)[2]
    # psi_2(R) = Psis(R)[:,2][1]*Phis(R)[1] + Psis(R)[:,2][2]*Phis(R)[2]
    # g_s_1 = zeros(201); g_s_2 = zeros(201); e_s_1 = zeros(201); e_s_2 = zeros(201); x = zeros(201);
    # for i in 0:0.1:20
    #     g_s_1[Int(i*10+1)] = psi_1(i)[1]
    #     g_s_2[Int(i*10+1)] = psi_1(i)[2]
    #     e_s_1[Int(i*10+1)] = psi_2(i)[1]
    #     e_s_2[Int(i*10+1)] = psi_2(i)[2]
    #     x[Int(i*10+1)] = i
    # end
    # scatter(x, g_s_1, markershape = :square, color = :blue, label = "ground_site_1")
    # scatter!(x, g_s_2, markershape = :square, color = :red, label = "ground_site_2")
    # scatter!(x, e_s_1, markershape = :diamond, color = :green, label = "excited_site_1")
    # scatter!(x, e_s_2, markershape = :diamond, color = :yellow, label = "excited_site_2")
    # gui()
    plot(r->psir_1(r), -R_0, R_0)
    plot!(r->psir_2(r)+1, -R_0, R_0)
    plot!(r->psil_1(r), -R_0, R_0)
    plot!(r->psil_2(r)+1, -R_0, R_0)
    gui()

end

"""
Non_adiabatic_states
N = number of sites
----------
Find a good estimate of the width of each state.
using exp(-(R-R_0)^2/2), A = 1.331335363800389
"""
function non_adiabatic_states(R,R_0s)
    phis = zeros(2)
    A = 1/0.354021770137867
    phi(R, R_0) = A*exp(-(R-R_0)^2/(0.01))
    for i in 1:length(R_0s)
        phis[i] = phi(R, R_0s[i])
    end
    return phis
end

"""
Overlap of diabatic states
"""
function overlap_phi(phi_1, phi_2)
    overlap_phi(R) = conj(phi_1(R))*phi_2(R)
    J_12 = riemann(overlap_phi, -10,10,1000)
    return J_12
end

"""
Function to calculate non-adiabatic coupling vector
using ADIABATIC states.
"""
function NACV(psi_i, psi_j, site)
    #grad_psi_1(R) = derivative(r->psi_1(r),R)
    grad_psi_j(R) = derivative(r->psi_j(site*r/2),R)
    overlap(R) = conj(psi_i(R))*grad_psi_j(R) #fix the time derivative of R.
    d_ij = riemann(overlap,-10,10,1000)
    # test(R) = conj(psi_i(R))*psi_i(R)
    # println(riemann(test,-10,10,1000))

    return d_ij

end

"""
Function to generate the full electronic wavefunction as a linear combination of
adiabatic_states by propagating the initial coeffients forwards in time (solving
the TDSE)
"""
function wavefunction_coefficients(d_12, d_21,a_i1, a_i2, a_p_m, a_p_p, R_e, dR, dt)
    # a_f1 = a_i1 - dt*im*a_p_m(R_e) - a_i2*NACV(psi_1, psi_2)*dR
    # a_f2 = a_i2 - dt*im*a_p_p(R_e) - a_i1*NACV(psi_2, psi_1)*dR
    Matrix = [-im*a_p_m(R_e) -d_12*dR/dt ; -d_21*dR/dt -im*a_p_p(R_e)]
    # println(M)
    lambdas = eigvals(Matrix)
    one,two,three,four = eigvecs(Matrix)
    C = (a_i1*four - a_i2*three)/(one*four - two*three)
    D = -(a_i1*two - a_i2*one)/(one*four - two*three)
    # println(C,"   ",D)
    # println(one,"    ", four)
    a_f1 = C*one*exp(lambdas[1]*dt) + D*three*exp(lambdas[2]*dt)
    a_f2 = C*two*exp(lambdas[1]*dt) + D*four*exp(lambdas[2]*dt)
    return a_f1, a_f2
end

"""
Funtion to generate the coefficients to link the diabatic and adiabatic states.
Psi = a1*psi_1 + a_2*psi_2
psi_1 = c_11*phi_1 + c_12*phi_2
psi_2 = c_21*phi_1 + c_22*phi_2

where psi_1 and psi_2 are the adiabatic states. for a 2 electron system these
are (0,1) and (1,0)

electron located at R_e; between 2 nuclei at R_n1 and R_n2
"""
function electronic_wavefunction(site, R_n, n_a_p_i, n_a_p_f, a_1, a_2, J_if, phis, plotting::Bool=true)
    R_e = site*R_n; R_n1 = -R_n; R_n2 = R_n
    H_diabatic = [n_a_p_i(R_e) J_if; J_if n_a_p_f(R_e)]
    phi_1(R) = phis(R)[1]; phi_2(R) = phis(R)[2]
    #H_adiabatic = [a_p_m(R_e) 0; 0 a_p_p(R_e)p3 = ]
    c_11, c_12, c_21, c_22 = eigvecs(H_diabatic)
    psi_1(R) = c_11*phi_1(R)[1] + c_12*phi_2(R)
    psi_2(R) = c_21*phi_1(R)[1] + c_22*phi_2(R)
    Psi(R) = a_1*psi_1(R) + a_2*psi_2(R)
    if plotting
        plot(r-> Psi(r), -10, 10)
    else
        return Psi, psi_1, psi_2
    end
end

"""
Function to calculate switching probabilities
Just considering state on V_11 to start.
"""
function g_12(d_12,dR, dt, a1, a2)
    g_12 = 2*real(conj(a2)*a1*dR*d_12/(dt*a1*conj(a1)))
    return g_12
end

"""
Function to calculate switching probabilities
Just considering state on V_22 to start.
"""
function g_21(d_21,dR, dt, a1, a2)
    g_21 = 2*real(conj(a1)*a2*dR*d_21/(dt*a2*conj(a2)))
    return g_21
end

"""
Function to define a change in parameters
"""
function r(R,Rn)
    r = R - Rn
    return r
end

"""
set up system parameters
plot initial state
propagate nuclei classically by dt
propagate wavefunction coefficients.

"""
function surface_hopping_2(R_0,x_0,v_0, n_a_p_i, n_a_p_f, a_1, a_2, e_i0, e_f0, phis)
# Initialise System
    PESm = true
    site = -1
    x1 = -x_0/2; x2 = x_0/2; x_n = x_0; v_n = v_0; PE = zeros(105000); KE = zeros(105000);TE = zeros(105000)
    # phis(R) = non_adiabatic_states(R, [x1, x2])
    phi_1(R) = phis(R)[1]
    phi_2(R) = phis(R)[2]
    J_if = overlap_phi(phi_1,phi_2)
    a_p_m = adiabatic_potential(K, R_0, e_i0, e_f0, J_if)[1]
    a_p_p = adiabatic_potential(K, R_0, e_i0, e_f0, J_if)[2]
    F_m, F_p = force_from_V(a_p_m, a_p_p, false)

    p1 = plot(r -> n_a_p_i(r),-R_0,R_0, color = :red)
    p1 = plot!(r -> n_a_p_f(r),-R_0,R_0, color = :red)
    p1 = plot!(r -> a_p_m(r),-R_0,R_0, color = :blue)
    p1 = plot!(r -> a_p_p(r),-R_0,R_0, color = :blue)
    p1 = scatter!([x_0/2],[a_p_m(x_0/2)], markershape = :circle, color = :black)
    p1 = scatter!([-x_0/2],[a_p_m(-x_0/2)], markershape = :circle, color = :black)
    gui()

    for j in 1:5000
        if PESm == true
            x,v = classical_propagation_2(real(x_n),real(v_n),0.01,0.0005,F_m)
        else
            x,v = classical_propagation_2(real(x_n),real(v_n),0.01,0.0005,F_p)
        end

        x1 = -x/2; x2 = x/2; Vg = [a_p_m(pos) for pos in x1];Ve = [a_p_p(pos) for pos in x1];

        KE[21*(j-1)+1:21*(j-1)+21] = 0.5*v.^2
        if PESm == true
            PE[21*(j-1)+1:21*(j-1)+21] = 2*Vg
            V_change = 2*(Ve-Vg)
        else
            PE[21*(j-1)+1:21*(j-1)+21] = 2*Ve
            V_change = 2*(Vg-Ve)
        end
        println("KE = ", KE[2*(j-1)+2], "PE = ", PE[2*(j-1)+2])
        TE[21*(j-1)+1:21*(j-1)+21] = KE[21*(j-1)+1:21*(j-1)+21]+PE[21*(j-1)+1:21*(j-1)+21]

        a1 = zeros(Complex,21); a2 = zeros(Complex,21); a1[1] = a_1; a2[1] = a_2
        x_n = x[end]; v_n = v[end]
        for i in 1:20
            phi_1n(R) = phi_1(R-(x_0/2-x2[i]))
            phi_2n(R) = phi_2(R+(x_0/2-x2[i]))
            phisn(R) = [phi_1n(R),phi_2n(R)]
            J_if = overlap_phi(phi_1n,phi_2n)
            a_p_m = adiabatic_potential(K, R_0, e_i0, e_f0, J_if)[1]
            a_p_p = adiabatic_potential(K, R_0, e_i0, e_f0, J_if)[2]
            F_m, F_p = force_from_V(a_p_m, a_p_p, false)
            Psi, psi_1, psi_2 = electronic_wavefunction(site, x2[i], n_a_p_i, n_a_p_f, a1[i], a2[i], J_if, phisn, false)
            d_12 = NACV(psi_1,psi_2, site); d_21 = NACV(psi_2, psi_1, site)
            a1[i+1], a2[i+1] = wavefunction_coefficients(d_12,d_21,a1[i],a2[i],a_p_m, a_p_p,site*x[i]/2,x[i+1]-x[i],0.0005)
        end
        Psi, psi_1, psi_2 = electronic_wavefunction(site, x2[end], n_a_p_i, n_a_p_f, a1[end], a2[end], J_if, phisn, false)

        if abs(Psi(x1[end])) > abs(Psi(x2[end]))
            site = -1
        else
            site = +1
        end

        if PESm == true
            d_12 = NACV(psi_1,psi_2,site)
            gm_12 = g_12(d_12,x[end]-x[end-1], 0.0005, a1[end], a2[end]); gp_21 = 0
        else
            d_21 = NACV(psi_2, psi_1,site)
            gp_21 = g_21(d_21,x[end]-x[end-1], 0.0005, a1[end], a2[end]); gm_12 =0
            println(gp_21)
        end

        chi = Distributions.Uniform()
        chi_rand = rand(chi)
        if chi_rand < 0.01
            println("small chi: ", chi_rand, "    ", max(real(gm_12), real(gp_21)))
        end
        if chi_rand < max(real(gm_12), real(gp_21))
            println(chi_rand, "    ", max(gm_12, gp_21))
            if PE[end] + KE[end] > V_change[end]
                v_0 = sqrt(complex(v[end]^2+2*(PE[end]-V_change[end])/M))

                if PESm
                    println(PESm," SWAP ", gm_12)
                    PESm = false
                    gm_12 = 0
                else
                    println(PESm," SWAP ", gp_21)
                    PESm = true
                    gp_21 = 0
                end
            end
        end
        #println(Psi(0.5))
        p1 = plot(r -> n_a_p_i(r),-R_0,R_0, color = :red)
        p1 = plot!(r -> n_a_p_f(r),-R_0,R_0, color = :red)
        p1 = plot!(r -> a_p_m(r),-R_0,R_0, color = :blue)
        p1 = plot!(r -> a_p_p(r),-R_0,R_0, color = :blue)
        p1 = plot!(r -> real(Psi(r))-1, -R_0, R_0, color = :cyan, linewidth = 4)
        p2 = scatter([x1[end]],[0], markershape = :circle, color = :black, xlims = (-R_0,R_0))
        p2 = scatter!([x2[end]],[0], markershape = :circle, color = :black)
        plot(p1,p2,layout = (2,1), size = (1200,800))
        gui()
    end
    TE = KE + PE
    plot(PE)

    gui()
end
# function test_wf_propagation()
#     plot_trajectory()
# end
"""
site = +/- 1 for nuclei on right or left

"""
function surface_hopping(T, dt, F_i, F_f, m, dx, R_i0, R_f0, n_a_p_i, n_a_p_f, J_if)

    PESm = true
    site = -1

    Total_E = zeros(T+2)
    xs = zeros(T+1)
    xs[1] = dx; x1_new = R_i0-dx/2; x2_new = R_f0 +dx/2; xe = x1_new; x_eq = R_f0-R_i0
    p1 = plot(r -> n_a_p_i(r),0,20)
    p1 = plot!(r -> n_a_p_f(r),0,20)
    t = 0; t_total = 0; gf_21 = 0; g1_12 = 0;
    x, v = classical_propagation_2(dx, 0., 100., 0.001)
    electronic_wavefunction = a_12

    while t+t_total<(T-2)
        #F(R) = F_m(R/2); M=1;
        x,v = classical_propagation(F_m, m, real(x), R_0, real(v), dt, slope_m)
        xs[Int(2+(t+t_total)/dt)] = real(x)
        xn_new = real(x)/2

        v_new = real(v)
        V1 = a_p_m(site*xn_new)
        V2 = a_p_p(site*xn_new)
        dR = (xs[Int(2+(t+t_total)/dt)] - xs[Int(1+(t+t_total)/dt)])
        if PESm
            Ve = V1; V_change = V2 - V1
            gi_12 = g_12(x2_new, xe, dR, dt, n_a_p_i, n_a_p_f, J_if)
            #println(x_new)
        else
            Ve = V2; V_change = V1 - V2
            gf_21 = g_21(x1_new, xe, dR, dt, n_a_p_i, n_a_p_f, J_if)
            #println(x_new)
        end
        t+=dt

        chi = Distributions.Uniform()
        chi_rand = rand(chi)
        if chi_rand < 0.01
            println(chi_rand, "    ", max(gi_12, gf_21))
        end
        K_E = 0.5*m*v_new^2
        P_E = Ve
        Total_E[Int(2+(t+t_total)/dt)] = K_E + P_E
        #println("chi = ", chi_rand, " g_12 = ", gi_12, " g_21 = ", gf_21)

        if chi_rand < max(gi_12, gf_21)

            if Total_E[Int(2+(t+t_total)/dt)] > V_change
                v_new2 = sqrt(v_new^2+2*(Ve-V_change)/m)

                if PESm
                    println(PESm," SWAP ", gi_12)
                    PESm = false
                    dx  = x2_new - R_f0
                    x, v = classical_propagation(F_f, m, dx, R_i0, R_f0, v_new2*sign(gi_12))
                    t_total += t
                    t = dt
                    gi_12 = 0
                else
                    println(PESm," SWAP ", gf_21)
                    PESm = true
                    dx  = x1_new - R_i0
                    x, v = classical_propagation(F_i, m, dx, R_i0, R_f0, v_new2*sign(gf_21))
                    t_total += t
                    t = dt
                    gf_21 = 0
                end
            end
        end

        p1 = plot(r -> n_a_p_i(r),0,20)
        p1 = plot!(r -> n_a_p_f(r),0,20)
        p1 = scatter!([x1_new],[V1], markershape = :circle, color = :black)
        p1 = scatter!([x2_new],[V2], markershape = :circle, color = :black)
        p1 = scatter!([xe],[Ve], markershape = :circle, color = :cyan, markersize = :3)
        gui()
    end

end

"""
Numeric integration, taken from http://mth229.github.io/integration.html
"""
function riemann(f::Function, a::Real, b::Real, n::Int; method="right")
    meth(f,l,r) = (1/6) * (f(l) + 4*(f((l+r)/2)) + f(r)) * (r-l)

  xs = a + (0:n) * (b-a)/n
  first = xs[1:end-1]; last = xs[2:end]; join = [first ; last];
  len = length(join)/2
  elements = reshape(join,(Int(len),2))
  pair = [elements[i,:] for i in 1:Int(len)]
  as = [meth(f,l,r) for (l,r) in pair]#meth(f, l, r)
  sum(as)

end


using Yao
using QuDiffEq
using LinearAlgebra
using OrdinaryDiffEq
using BitBasis
using Random
using Test
using YaoBlocks

#Linear Diff Equation Unitary M
function diffeqProblem(nbit::Int)
    siz = 1<<nbit
    Au = rand_unitary(siz)
    An = rand(ComplexF64,siz,siz)
    b = normalize!(rand(ComplexF64, siz))
    x = normalize!(rand(ComplexF64, siz))
    Au,An,b, x
end

@testset "QuLDE_Test" begin
    Random.seed!(2)
    N = 1
    k = 3
    tspan = (0.0,0.4)
    Au,An,b,x = diffeqProblem(N)
    qprob = QuLDEProblem(Au, b, x, tspan)
    f(u,p,t) = Au*u + b;
    prob = ODEProblem(f, x, tspan)

    sol = solve(prob, Tsit5(), dt = 0.1, adaptive = false)
    s = sol.u[end]

    out = solve(qprob, QuLDE(k))

    @test isapprox.(s, out, atol = 0.01) |> all

    qprob = QuLDEProblem(An, b, x, tspan)
    f(u,p,t) = An*u + b;
    prob = ODEProblem(f, x, tspan)

    sol = solve(prob, Tsit5(), dt = 0.1, adaptive = false)
    s = sol.u[end]

    out = solve(qprob, QuLDE(k))
    @test isapprox.(s, out, atol = 0.02) |> all

    # u0 equal to zero
    tspan = (0.0,0.1)

    qprob = QuLDEProblem(Au,b,tspan)
    out = solve(qprob,QuLDE(k))
    t = tspan[2] - tspan[1]
    r_out = (exp(Au*t) - Diagonal(ones(length(b))))*Au^(-1)*b
    @test isapprox.(r_out, out, atol = 1e-3) |> all

    qprob = QuLDEProblem(An,b,tspan)
    out = solve(qprob,QuLDE(k))
    t = tspan[2] - tspan[1]
    r_out = (exp(An*t) - Diagonal(ones(length(b))))*An^(-1)*b
    @test isapprox.(r_out, out, atol = 1e-3) |> all
end


doc"""
    InverseGamma(α, θ)

The *inverse gamma distribution* with shape parameter `α` and scale `θ` has probability
density function

$f(x; \alpha, \theta) = \frac{\theta^\alpha x^{-(\alpha + 1)}}{\Gamma(\alpha)}
e^{-\frac{\theta}{x}}, \quad x > 0$

It is related to the [`Gamma`](:func:`Gamma`) distribution: if $X \sim \operatorname{Gamma}(\alpha, \beta)$, then $1 / X \sim \operatorname{InverseGamma}(\alpha, \beta^{-1})$.

.. code-block:: julia

    InverseGamma()        # Inverse Gamma distribution with unit shape and unit scale, i.e. InverseGamma(1, 1)
    InverseGamma(a)       # Inverse Gamma distribution with shape a and unit scale, i.e. InverseGamma(a, 1)
    InverseGamma(a, b)    # Inverse Gamma distribution with shape a and scale b

    params(d)        # Get the parameters, i.e. (a, b)
    shape(d)         # Get the shape parameter, i.e. a
    scale(d)         # Get the scale parameter, i.e. b

External links

* [Inverse gamma distribution on Wikipedia](http://en.wikipedia.org/wiki/Inverse-gamma_distribution)

"""
immutable InverseGamma{T<:Real} <: ContinuousUnivariateDistribution
    invd::Gamma{T}
    θ::T

    function InverseGamma(α, θ)
        @check_args(InverseGamma, α > zero(α) && θ > zero(θ))
        new(Gamma(α, 1 / θ), θ)
    end
end

InverseGamma{T<:Real}(α::T, θ::T) = InverseGamma{T}(α, θ)
InverseGamma(α::Real, θ::Real) = InverseGamma(promote(α, θ)...)
InverseGamma(α::Integer, θ::Integer) = InverseGamma(Float64(α), Float64(θ))
InverseGamma(α::Real) = InverseGamma(α, 1.0)
InverseGamma() = InverseGamma(1.0, 1.0)

@distr_support InverseGamma 0.0 Inf

#### Conversions
convert{T <: Real, S <: Real}(::Type{InverseGamma{T}}, α::S, θ::S) = InverseGamma(T(α), T(θ))
convert{T <: Real, S <: Real}(::Type{InverseGamma{T}}, d::InverseGamma{S}) = InverseGamma(T(shape(d.invd)), T(d.θ))

#### Parameters

shape(d::InverseGamma) = shape(d.invd)
scale(d::InverseGamma) = d.θ
rate(d::InverseGamma) = scale(d.invd)

params(d::InverseGamma) = (shape(d), scale(d))
@inline partype{T<:Real}(d::InverseGamma{T}) = T

#### Parameters

mean{T<:Real}(d::InverseGamma{T}) = ((α, θ) = params(d); α  > 1 ? θ / (α - 1) : T(Inf))

mode(d::InverseGamma) = scale(d) / (shape(d) + 1)

function var{T<:Real}(d::InverseGamma{T})
    (α, θ) = params(d)
    α > 2 ? θ^2 / ((α - 1)^2 * (α - 2)) : T(Inf)
end

function skewness{T<:Real}(d::InverseGamma{T})
    α = shape(d)
    α > 3 ? 4sqrt(α - 2) / (α - 3) : T(NaN)
end

function kurtosis{T<:Real}(d::InverseGamma{T})
    α = shape(d)
    α > 4 ? (30α - 66) / ((α - 3) * (α - 4)) : T(NaN)
end

function entropy(d::InverseGamma)
    (α, θ) = params(d)
    α + lgamma(α) - (1 + α) * digamma(α) + log(θ)
end

#### Evaluation

pdf(d::InverseGamma, x::Real) = exp(logpdf(d, x))

function logpdf(d::InverseGamma, x::Real)
    (α, θ) = params(d)
    α * log(θ) - lgamma(α) - (α + 1) * log(x) - θ / x
end

cdf(d::InverseGamma, x::Real) = ccdf(d.invd, 1 / x)
ccdf(d::InverseGamma, x::Real) = cdf(d.invd, 1 / x)
logcdf(d::InverseGamma, x::Real) = logccdf(d.invd, 1 / x)
logccdf(d::InverseGamma, x::Real) = logcdf(d.invd, 1 / x)

quantile(d::InverseGamma, p::Real) = 1 / cquantile(d.invd, p)
cquantile(d::InverseGamma, p::Real) = 1 / quantile(d.invd, p)
invlogcdf(d::InverseGamma, p::Real) = 1 / invlogccdf(d.invd, p)
invlogccdf(d::InverseGamma, p::Real) = 1 / invlogcdf(d.invd, p)

function mgf{T<:Real}(d::InverseGamma{T}, t::Real)
    (a, b) = params(d)
    t == zero(t) ? one(T) : 2(-b*t)^(0.5a) / gamma(a) * besselk(a, sqrt(-4*b*t))
end

function cf{T<:Real}(d::InverseGamma{T}, t::Real)
    (a, b) = params(d)
    t == zero(t) ? one(T)+zero(T)*im : 2(-im*b*t)^(0.5a) / gamma(a) * besselk(a, sqrt(-4*im*b*t))
end

#### Evaluation

rand(d::InverseGamma) = 1 / rand(d.invd)

function _rand!(d::InverseGamma, A::AbstractArray)
    s = sampler(d.invd)
    for i = 1:length(A)
    	v = 1 / rand(s)
        @inbounds A[i] = v
    end
    A
end


@testset "Diagonal" begin

    # diag:
    let P = 10, rng = MersenneTwister(123456), N = 10, k = 2

        # Some rngs.
        vP, mPP, vPk = ()->randn(rng, P), ()->randn(rng, P, P), ()->randn(rng, P - k)

        # Test on-central-diagonal `diag`.
        X0 = randn!(rng, similar(mPP()))
        @test check_errs(N, unary_ȲD(diag)..., vP, mPP, mPP)
        @test check_errs(N, unary_ȲD_inplace(diag, X0)..., vP, mPP, mPP)

        # Test off-central-diagonal `diag`.
        _diag = X->diag(X, k)
        _∇diag = (y, ȳ, X)->∇(diag, Val{1}, (), y, ȳ, X, k)
        _∇diag_inp = (y, ȳ, X)->∇(copy(X0), diag, Val{1}, (), y, ȳ, X, k) - X0
        @test check_errs(N, _diag, _∇diag, vPk, mPP, mPP)
        @test check_errs(N, _diag, _∇diag_inp, vPk, mPP, mPP)
    end

    # diagm:
    let P = 10, rng = MersenneTwister(123456), N = 10, k = 3

        # Some rngs.
        sc, vP = ()->randn(rng), ()->randn(rng, P)
        mPP, mPPk = ()->randn(rng, P, P), ()->randn(rng, P + k, P + k)

        # Test on-central-diagonal `diagm`.
        x0 = randn!(rng, similar(vP()))
        @test check_errs(N, unary_ȲD(diagm)..., mPP, vP, vP)
        @test check_errs(N, unary_ȲD_inplace(diagm, x0)..., mPP, vP, vP)

        # Test off-central-diagonal `diagm`.
        _diagm = x->diagm(x, k)
        _∇diagm = (Y, Ȳ, x)->∇(diagm, Val{1}, (), Y, Ȳ, x, k)
        _∇diagm_inp = (Y, Ȳ, x)->∇(copy(x0), diagm, Val{1}, (), Y, Ȳ, x, k) - x0
        @test check_errs(N, _diagm, _∇diagm, mPPk, vP, vP)
        @test check_errs(N, _diagm, _∇diagm_inp, mPPk, vP, vP)

        # Test scalar input to `diagm`.
        @test check_errs(N, unary_ȲD(diagm)..., ()->randn(rng, 1, 1), sc, sc)
    end

    # Diagonal:
    let P = 10, rng = MersenneTwister(123456), N = 10
        sc, vP = ()->abs(randn(rng)), ()->abs.(randn(rng, P))
        mPP, dPP = ()->abs.(randn(rng, P, P)), ()->Diagonal(abs.(randn(rng, P)))
        D0 = Diagonal(randn!(rng, similar(vP())))

        # Construction.
        @test check_errs(N, unary_ȲD(Diagonal)..., dPP, vP, vP)
        @test check_errs(N, unary_ȲD_inplace(Diagonal, vP())..., dPP, vP, vP)
        @test check_errs(N, unary_ȲD(Diagonal)..., dPP, mPP, mPP)
        @test check_errs(N, unary_ȲD_inplace(Diagonal, mPP())..., dPP, mPP, mPP)

        # Determinant.
        @test check_errs(N, unary_ȲD(det)..., sc, dPP, dPP)
        @test check_errs(N, unary_ȲD_inplace(det, D0)..., sc, dPP, dPP)

        # Log Determinant.
        @test check_errs(N, unary_ȲD(logdet)..., sc, dPP, dPP)
        @test check_errs(N, unary_ȲD_inplace(logdet, D0)..., sc, dPP, dPP)
    end
end


using ModelingToolkit
using Test

# `Expr`, `Number` -> `Operation`
@Param a
@Param b
@Unknown x(t) y()
@test convert(Expression, 2) == 2
expr = :(-inv(2sqrt(+($a, $b))))
op   = Operation(-, [Operation(inv,
                [Operation(*, [2, Operation(sqrt,
                              [Operation(+, [a, b])])])])])
@test convert(Expression, expr) == op
expr1 = :($x^($y-1))
op1   = Operation(^, [x, Operation(-, [y, 1])])
@test convert(Expression, expr1) == op1

@test_throws ArgumentError convert(Expression, :([a, b]))
@test_throws ArgumentError convert(Expression, :(a ? b : c))


"""A face of Box{3}"""
struct Rectangle
    # Origin
    O::Point{3}
    # Two perp vector defining the plane ...
    t1::Point{3}
    t2::Point{3}
    # ... with normal
    normal::Point{3}
    # Lengths of sides
    l1::Float64
    l2::Float64
    # A
    # |
    # O--B
    function Rectangle(O::Point{3}, t1::Point{3}, t2::Point{3}, n::Point{3})
        l1 = norm(t1)
        t1 /= l1

        l2 = norm(t2)
        t2 /= l2

        n0 = cross(t1, t2)
        @assert n == normalize(n0)
        new(O, t1, t2, n, l1, l2)
    end
end

"""Construct from enumerating"""
function Rectangle(b::Box{3}, facet::Int)
    const dX = b.Y - b.X

    facet == 1 && return Rectangle(b.X, Point(dX[1], 0, 0), Point(0, dX[2], 0), Point(0, 0, 1))
    facet == 2 && return Rectangle(b.X, Point(0, dX[2], 0), Point(0, 0, dX[3]), Point(1, 0, 0))
    facet == 3 && return Rectangle(b.X, Point(0, 0, dX[3]), Point(dX[1], 0, 0), Point(0, 1, 0))

    facet == 4 && return Rectangle(b.Y, Point(-dX[1], 0, 0), Point(0, -dX[2], 0), Point(0, 0, 1))
    facet == 5 && return Rectangle(b.Y, Point(0, -dX[2], 0), Point(0, 0, -dX[3]), Point(1, 0, 0))
    facet == 6 && return Rectangle(b.Y, Point(0, 0, -dX[3]), Point(-dX[1], 0, 0), Point(0, 1, 0))

end

# Eval the curve using [0, 1]^2 parametrization
(r::Rectangle)(s::Number, t::Number) = r.O + s*r.l1*r.t1 + t*r.l2*r.t2

show(io::IO, b::Rectangle) = print(io, "$(b.O)-$(b(1, 1))")

"""
    Quad tree like division

    Bottom part    Top part
    |---|---|      |---|---|
    | 2 | 4 |      | 6 | 8 |
    |---|---|      |---|---|
    | 1 | 3 |      | 5 | 7 |
    |---|---|      |---|---|

"""
function split(box::Box{3})
    X, Y = box.X, box.Y
    M = 0.5(X + Y)
    [Box(Point(X[1], X[2], X[2]), Point(M[1], M[2], M[2])),
     Box(Point(X[1], M[2], X[2]), Point(M[1], Y[2], M[2])),
     Box(Point(M[1], X[2], X[2]), Point(Y[1], M[2], M[2])),
     Box(Point(M[1], M[2], X[2]), Point(Y[1], Y[2], M[2])),
     Box(Point(X[1], X[2], M[2]), Point(M[1], M[2], Y[2])),
     Box(Point(X[1], M[2], M[2]), Point(M[1], Y[2], Y[2])),
     Box(Point(M[1], X[2], M[2]), Point(Y[1], M[2], Y[2])),
     Box(Point(M[1], M[2], M[2]), Point(Y[1], Y[2], Y[2]))]
end

"""6 faces of the prism"""
function surface(box::Box{3})
    map(i -> Rectangle(box, i), 1:6)
end

"""Does a face(rectangle) collide with segment"""
function collides(line::Segment{3}, box::Rectangle)
    D = line.O - box.O
    D1 = dot(D, box.t1)
    D2 = dot(D, box.t2)
    D3 = dot(D, line.t)

    c1 = dot(box.t1, line.t)
    c2 = dot(box.t2, line.t)
    C = c1^2 + c2^2 - 1

    # @show [1 0 -c1; 0 1 -c2; c1 c2 -1]\[D1, D2, D3]
    # We solve the 3x3 system
    α = (-c1*c1/C + 1)*D1 + (-c1*c2/C)*D2 + (c1/C)*D3
    # @show α
    !(-EPS < α < box.l1 + EPS) && return false

    β = (-c1*c2/C)*D1 + (-c2*c2/C + 1)*D2 + (c2/C)*D3
    # @show β
    !(-EPS < β < box.l2 + EPS) && return false

    γ = (-c1/C)*D1 + (-c2/C)*D2 + (1/C)*D3
    # @show γ
    !(-EPS < γ < line.length + EPS) && return false

    true
end

"""See if segment collides with a box"""
function collides(segment::Segment{3}, box::Box{3})
    !bbox_collides(segment.box, box) && return false

    (segment.O ∈ box && segment.O ∈ box) && return true

    # Check collisions with cases --[---]--
    for line in surface(box)
        collides(segment, line) && return true
    end

    false
end