function lp1_test()
    #=
        min -4x_2 - 1
    s.t.
        x_1 + 2x_2 <= 3
        x_1 >= 3
    =#
    model = TestModel{Float64}()

    X = MOI.add_variables(model, 2)

    MOI.add_constraint(model,
        MOI.ScalarAffineFunction(
            MOI.ScalarAffineTerm.([1.0, 2.0], X), 0.0),
            MOI.LessThan(3.0))

    MOI.add_constraint(model,
        MOI.SingleVariable(X[1]),
            MOI.GreaterThan(3.0))

    MOI.set(model,
        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(),
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([-4.0], [X[2]]), -1.0))

    MOI.set(model, MOI.ObjectiveSense(), MOI.MIN_SENSE)

    return model
end

function lp2_test()
    #=
        min -4x1 -3x2 -1
    s.t.
        2x1 + x2 + 1 <= 4
        x1 + 2x2 + 1 <= 4
        x1 >= 1
        x2 >= 0
    =#
    model = TestModel{Float64}()

    X = MOI.add_variables(model, 2)

    MOI.add_constraint(model,
        MOI.ScalarAffineFunction(
            MOI.ScalarAffineTerm.([2.0, 1.0], X), 0.0),
             MOI.LessThan(3.0))

    MOI.add_constraint(model,
        MOI.ScalarAffineFunction(
            MOI.ScalarAffineTerm.([1.0, 2.0], X), 0.0),
             MOI.LessThan(3.0))

    MOI.add_constraint(model,
        MOI.ScalarAffineFunction(
            [MOI.ScalarAffineTerm(1.0, X[1])], 0.0),
             MOI.GreaterThan(1.0))

    MOI.add_constraint(model,
        MOI.ScalarAffineFunction(
            [MOI.ScalarAffineTerm(1.0, X[2])], 0.0),
             MOI.GreaterThan(0.0))

    MOI.set(model,
        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(),
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([-4.0, -3.0], X), -1.0))

    MOI.set(model, MOI.ObjectiveSense(), MOI.MIN_SENSE)

    return model
end

function lp3_test()
    #=
        min -4x1 -3x2 -1
    s.t.
        2x1 + x2 + 1 <= 4
        x1 + 2x2 + 1 <= 4
        x1 >= 1
        x2 >= 0
    =#
    model = TestModel{Float64}()

    X = MOI.add_variables(model, 2)

    MOI.add_constraint(model,
        MOI.ScalarAffineFunction(
            MOI.ScalarAffineTerm.([2.0, 1.0], X), 0.0),
             MOI.LessThan(3.0))

    MOI.add_constraint(model,
        MOI.ScalarAffineFunction(
            MOI.ScalarAffineTerm.([1.0, 2.0], X), 0.0),
             MOI.LessThan(3.0))

    MOI.add_constraint(model,
        MOI.SingleVariable(X[1]),
             MOI.GreaterThan(1.0))

    MOI.add_constraint(model,
        MOI.SingleVariable(X[2]),
             MOI.GreaterThan(0.0))

    MOI.set(model,
        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(),
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([-4.0, -3.0], X), -1.0))

    MOI.set(model, MOI.ObjectiveSense(), MOI.MIN_SENSE)

    return model
end

function lp4_test()
    #=
        max 4x1 + 3x2
      s.t.
        x1 >= 1
        x2 >= 0
    =#
    model = TestModel{Float64}()

    X = MOI.add_variables(model, 2)

    MOI.add_constraint(model,
        MOI.SingleVariable(X[1]),
             MOI.GreaterThan(1.0))

    MOI.add_constraint(model,
        MOI.SingleVariable(X[2]),
             MOI.GreaterThan(0.0))

    MOI.set(model,
        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(),
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([4.0, 3.0], X), 0.0))

    MOI.set(model, MOI.ObjectiveSense(), MOI.MAX_SENSE)

    return model
end

function lp5_test()
    #=
        min -4x1 -3x2 -1
    s.t.
        2x1 + x2  == 3
        x1 + 2x2  == 3
        x1 >= 1
        x2 == 0
    =#
    model= TestModel{Float64}()

    X = MOI.add_variables(model, 2)

    MOI.add_constraint(model,
        MOI.ScalarAffineFunction(
            MOI.ScalarAffineTerm.([2.0, 1.0], X), 0.0),
            MOI.EqualTo(3.0))

    MOI.add_constraint(model,
        MOI.ScalarAffineFunction(
            MOI.ScalarAffineTerm.([1.0, 2.0], X), 0.0),
            MOI.EqualTo(3.0))

    MOI.add_constraint(model,
        MOI.SingleVariable(X[1]),
            MOI.GreaterThan(1.0))

    MOI.add_constraint(model,
        MOI.SingleVariable(X[2]),
            MOI.EqualTo(0.0))

    MOI.set(model,
        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(),
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([-4.0, -3.0], X), -1.0))

    MOI.set(model, MOI.ObjectiveSense(), MOI.MIN_SENSE)

    return model
end

function lp6_test()
    #=
        min -4x1 -3x2 -1
    s.t.
        2x1 + x2 - 3 <= 0
        x1 + 2x2 - 3 <= 0
        x1 >= 1
        x2 >= 0
    =#
    model = TestModel{Float64}()

    X = MOI.add_variables(model, 2)

    c1 = MOI.VectorAffineFunction(
                MOI.VectorAffineTerm.([1, 1, 2, 2], MOI.ScalarAffineTerm.([2.0, 1.0, 1.0, 2.0], [X; X])),
                [-3.0, -3.0])

    MOI.add_constraint(model, c1, MOI.Nonpositives(2))

    MOI.add_constraint(model,
        MOI.SingleVariable(X[1]),
            MOI.GreaterThan(1.0))

    MOI.add_constraint(model,
        MOI.SingleVariable(X[2]),
            MOI.GreaterThan(0.0))

    MOI.set(model,
        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(),
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([-4.0, -3.0], X), -1.0))

    MOI.set(model, MOI.ObjectiveSense(), MOI.MIN_SENSE)

    return model
end

function lp7_test()
    #=
        min -4x1 -3x2 -1
    s.t.
        2x1 + x2 - 3 <= 0
        x1 + 2x2 - 3 <= 0
        x1 >= 1
        x2 >= 0
    =#
    model= TestModel{Float64}()

    X = MOI.add_variables(model, 2)

    c1 = MOI.VectorAffineFunction(
            MOI.VectorAffineTerm.([1, 1, 2, 2], MOI.ScalarAffineTerm.([2.0, 1.0, 1.0, 2.0], [X; X])),
            [-3.0, -3.0])

    MOI.add_constraint(model, c1, MOI.Nonpositives(2))

    MOI.add_constraint(model,
        MOI.SingleVariable(X[1]),
            MOI.GreaterThan(1.0))

    MOI.add_constraint(model,
        MOI.SingleVariable(X[2]),
            MOI.GreaterThan(0.0))

    MOI.set(model,
        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(),
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([-4.0, -3.0], X), -1.0))

    MOI.set(model, MOI.ObjectiveSense(), MOI.MIN_SENSE)

    return model
end

function lp8_test()
    #=
        min -4x1 -3x2 -1
    s.t.
        2x1 + x2 - 3 <= 0
        x1 + 2x2 - 3 <= 0
        x1 >= 1
        x2 >= 0
    =#
    model = TestModel{Int64}()

    X = MOI.add_variables(model, 2)

    c1 = MOI.VectorAffineFunction(
            MOI.VectorAffineTerm.([1, 1, 2, 2], MOI.ScalarAffineTerm.([2, 1, 1, 2], [X; X])),
            [-3, -3])

    MOI.add_constraint(model, c1, MOI.Nonpositives(2))

    c2 = MOI.VectorAffineFunction(
            MOI.VectorAffineTerm.([1, 2], MOI.ScalarAffineTerm.([1, 1], X)),
            [1, 0])

    MOI.add_constraint(model, c2, MOI.Nonnegatives(2))

    MOI.set(model,
        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(),
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([-4, -3], X), -1))

    MOI.set(model, MOI.ObjectiveSense(), MOI.MIN_SENSE)

    return model
end

function lp9_test()
    #=
        min x + y
    s.t.
        -1 <= x + y <= 10
        x,  y >= 0
    =#
    model = TestModel{Float64}()

    x = MOI.add_variable(model)
    y = MOI.add_variable(model)

    vc = MOI.add_constraints(model,
        [MOI.SingleVariable(x), MOI.SingleVariable(y)],
        [MOI.GreaterThan(0.0), MOI.GreaterThan(0.0)])

    c = MOI.add_constraint(model, MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([1.0, 1.0], [x,y]), 0.0), MOI.Interval(-1.0, 10.0))

    MOI.set(model,
        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(),
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([1.0, 1.0], [x, y]), 0.0))

    MOI.set(model, MOI.ObjectiveSense(), MOI.MIN_SENSE)

    return model
end

function lp10_test()
    #=
        min x1
    s.t.
        2x1 + x2  == 3
        x1 + 2x2  == 3
        x1 >= 1
        x2 == 0
    =#
    model= TestModel{Float64}()

    X = MOI.add_variables(model, 2)

    MOI.add_constraint(model,
        MOI.ScalarAffineFunction(
            MOI.ScalarAffineTerm.([2.0, 1.0], X), 0.0),
            MOI.EqualTo(3.0))

    MOI.add_constraint(model,
        MOI.ScalarAffineFunction(
            MOI.ScalarAffineTerm.([1.0, 2.0], X), 0.0),
            MOI.EqualTo(3.0))

    MOI.add_constraint(model,
        MOI.SingleVariable(X[1]),
            MOI.GreaterThan(1.0))

    MOI.add_constraint(model,
        MOI.SingleVariable(X[2]),
            MOI.EqualTo(0.0))

    MOI.set(model,
        MOI.ObjectiveFunction{MOI.SingleVariable}(),
        MOI.SingleVariable(X[1]))

    MOI.set(model, MOI.ObjectiveSense(), MOI.MIN_SENSE)

    return model
end

function lp11_test()
    #=
        Feasibility
    s.t.
        x1 >= 1
        x2 >= 0
    =#
    model = TestModel{Float64}()

    X = MOI.add_variables(model, 2)

    MOI.add_constraint(model,
        MOI.SingleVariable(X[1]),
             MOI.GreaterThan(1.0))

    MOI.add_constraint(model,
        MOI.SingleVariable(X[2]),
             MOI.GreaterThan(0.0))

    MOI.set(model,
        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(),
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([4.0, 3.0], X), 0.0))

    return model
end

function lp12_test()
    #=
        min 4x_3 + 5
    s.t.
        x_1 + 2x_2 + x_3 <= 20
        x_1 <= 1
        x_2 <= 3
    =#
    model = TestModel{Float64}()

    X = MOI.add_variables(model, 3)

    MOI.add_constraint(model,
        MOI.ScalarAffineFunction(
            MOI.ScalarAffineTerm.([1.0, 2.0, 1.0], X), 0.0),
            MOI.LessThan(20.0))

    MOI.add_constraint(model,
        MOI.SingleVariable(X[1]),
            MOI.LessThan(1.0))

    MOI.add_constraint(model,
        MOI.SingleVariable(X[2]),
            MOI.LessThan(3.0))

    MOI.set(model,
        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(),
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([4.0], [X[3]]), 5.0))

    MOI.set(model, MOI.ObjectiveSense(), MOI.MIN_SENSE)

    return model
end

function lp13_test()
    #=
        min -4x1 -3x2 -1
    s.t.
        2x1 + x2 - 3 <= 0
        x1 + 2x2 - 3 <= 0
        x1 >= 0
        x2 >= 0
    =#
    model= TestModel{Float64}()

    X = MOI.add_variables(model, 2)

    c1 = MOI.VectorAffineFunction(
            MOI.VectorAffineTerm.([1, 1, 2, 2], MOI.ScalarAffineTerm.([2.0, 1.0, 1.0, 2.0], [X; X])),
            [-3.0, -3.0])

    MOI.add_constraint(model, c1, MOI.Nonpositives(2))

    MOI.add_constraint(model, MOI.VectorOfVariables(X), MOI.Nonnegatives(2))

    MOI.set(model,
        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(),
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([-4.0, -3.0], X), -1.0))

    MOI.set(model, MOI.ObjectiveSense(), MOI.MIN_SENSE)

    return model
end


#=

  Discrete tomography in Julia ConstraintSolver.jl

  Note: The origin of the problem is from ECLiPSe,
  but this model has been transformed in this way
     MiniZinc -> SICStus Prolog -> ECLiPSe -> B-Prolog -> Picat -> ConstraintSolver.jl
  Here is my own take at the problem.

  Problem from http://eclipse-clp.org/examples/tomo.ecl.txt
  """
  This is a little "tomography" problem, taken from an old issue
  of Scientific American.

  A matrix which contains zeroes and ones gets "x-rayed" vertically and
  horizontally, giving the total number of ones in each row and column.
  The problem is to reconstruct the contents of the matrix from this
  information. Sample run:

  ?- go.
     0 0 7 1 6 3 4 5 2 7 0 0
  0
  0
  8      * * * * * * * *
  2      *             *
  6      *   * * * *   *
  4      *   *     *   *
  5      *   *   * *   *
  3      *   *         *
  7      *   * * * * * *
  0
  0

  Eclipse solution by Joachim Schimpf, IC-Parc
  """

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

=#

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

#
# print_matrix_str(x,zstring)
#
# Print each row and convert each element to the corresponding position
# in zstring. "z" is for zero-based, i.e. first char corresponds to 0, etc
#
function print_matrix_str(x,zstring="")
    for row in eachrow(x)
        if zstring == ""
            println(row)
        else
            println(join(row.|>c->zstring[c+1]))
        end
    end
    println()
end

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

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

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

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

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

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

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

                                                            "time_limit"=>6,

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

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

    row_sums = problem[:r]
    col_sums = problem[:c]

    rows = length(row_sums)
    cols = length(col_sums)

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

    # Check rows and cols
    for i in 1:rows
        @constraint(model,row_sums[i] == sum(x[i,:]))
    end
    for j in 1:cols
        @constraint(model,col_sums[j] == sum(x[:,j]))
    end

    # Solve the problem
    optimize!(model)

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

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

    return status
end

discrete_tomography_problems = Dict(

#
# The three first problems are from the ECLiPSe model:
#
# The above stated problem
:1  => Dict(
       :r => [0,0,8,2,6,4,5,3,7,0,0], # row sums
       :c => [0,0,7,1,6,3,4,5,2,7,0,0] # column sums
       ),

:2 => Dict(
       :r => [10,4,8,5,6],
       :c => [5,3,4,0,5,0,5,2,2,0,1,5,1]
),

# This give three slightly different solutions.
:3 => Dict(
        :r => [11,5,4],
        :c => [3,2,3,1,1,1,1,2,3,2,1]
),

# This is my own problem.
:4 => Dict(
        :r => [0,2,2,2,2,2,8,8,4,4,4,4,4,0],
        :c => [0,0,0,12,12,2,2,2,2,7,7,0,0,0]
),

)

for p in sort(collect(keys(discrete_tomography_problems)))
    println("problem $p")
    @time discrete_tomography(discrete_tomography_problems[p])
end


@testset "GHC" begin
  Z = [ones(10,10) 2ones(10,10); 3ones(10,10) 4ones(10,10)]
  𝒮 = georef((Z=Z,))
  p = partition(𝒮, GHC(4, 1.0))
  @test length(p) == 4
  @test all(nelements.(p) .== 100)
  p1, p2, p3, p4 = p
  @test mean(coordinates(centroid(p1, ind)) for ind in 1:nelements(p1)) == [5.0,5.0]
  @test mean(coordinates(centroid(p2, ind)) for ind in 1:nelements(p2)) == [15.0,5.0]
  @test mean(coordinates(centroid(p3, ind)) for ind in 1:nelements(p3)) == [5.0,15.0]
  @test mean(coordinates(centroid(p4, ind)) for ind in 1:nelements(p4)) == [15.0,15.0]

  C = cluster(𝒮, GHC(4, 1.0))
  @test C.cluster == vec(Z')

  𝒮′ = georef(values(𝒮), centroid.(domain(𝒮)))
  C′ = cluster(𝒮′, GHC(4, 1.0))
  @test C.cluster == C′.cluster

  𝒮 = georef((z=[√(i^2+j^2) for i in 1:50, j in 1:50],))
  p = partition(𝒮, GHC(50, 1.0))
  @test length(p) == 50

  if visualtests
    @test_reference "data/ghc.png" plot(p)
  end
end


using MultivariateTests
using Compat
using Compat.Test
using Compat.DelimitedFiles
using StatsBase
using HypothesisTests

# Columns are calcium, iron, protein, vitamin A, vitamin C
nutrient = readdlm(joinpath(@__DIR__, "data", "nutrient.txt"))[:,2:end]

# Columns are information, similarities, arithmetic, picture completion
wechsler = readdlm(joinpath(@__DIR__, "data", "wechsler.txt"))[:,2:end]

# Columns are type, length, left, right, bottom, top, diag
swiss = readdlm(joinpath(@__DIR__, "data", "swiss3.txt"))
genuine = convert(Matrix{Float64}, swiss[view(swiss, :, 1) .== "real", 2:end])
counterfeit = convert(Matrix{Float64}, swiss[view(swiss, :, 1) .== "fake", 2:end])

# Columns are survey answers: husbands' to questions 1-4, wives' to questions 1-4
spouse = readdlm(joinpath(@__DIR__, "data", "spouse.txt"))

@testset "Utility functions" begin
    MT = MultivariateTests

    @test MT.checkdims(rand(3, 2), rand(4, 2)) == (2, 3, 4)
    # Mismatched number of variables
    @test_throws DimensionMismatch MT.checkdims(rand(4, 3), rand(4, 6))
    # Empty observations
    @test_throws ArgumentError MT.checkdims(rand(0, 4), rand(2, 4))

    Sx = [ 3.0 -1.5 0.0
          -1.5  1.0 0.5
           0.0  0.5 1.0]
    Sy = fill(4.5, (3, 3))
    out = [3.5 0.5      1.5
           0.5 2.166667 1.833333
           1.5 1.833333 2.166667]
    @test MT._poolcov!(Sx, 2, Sy, 1) ≈ out atol=1e-6
    # Input is modified
    @test Sx ≈ out atol=1e-6

    # Positive semi-definite but not positive definite
    P = [ 1.0000  0.7426  0.1601 -0.7000 0.5500
          0.7426  1.0000 -0.2133 -0.5818 0.5000
          0.1601 -0.2133  1.0000 -0.1121 0.1000
         -0.7000 -0.5818 -0.1121  1.0000 0.4500
          0.5500  0.5000  0.1000  0.4500 1.0000]
    # Positive definite
    Q = [ 1.0 -0.5  0.0
         -0.5  1.0 -0.5
          0.0 -0.5  1.0]
    @test @inferred(MT.At_Binv_A(ones(5), P)) ≈ -0.8008792 atol=1e-6
    @test @inferred(MT.At_Binv_A(ones(3), Q)) ≈ 10.0
end

@testset "Generalizations of sample variance" begin
    @test @inferred(genvar(nutrient)) ≈ 2.8310418e19 atol=1e-6
    @test @inferred(totalvar(nutrient)) ≈ 2.83266877e6 atol=1e-6

    X = [1 2 5
         4 1 6
         4 0 4]
    @test @inferred(genvar(X)) ≈ 0.0
    @test @inferred(totalvar(X)) ≈ 5.0

    x = rand(Float32, 10)
    @test genvar(x) == totalvar(x) == var(x)

    it = (xᵢ for xᵢ in x)
    @test genvar(it) == totalvar(it) == var(it)
end

@testset "Partial correlation" begin
    @test @inferred(partialcor(wechsler[:,1], wechsler[:,2], wechsler[:,3:4])) ≈ 0.7118787 atol=1e-6

    w = PartialCorTest(wechsler[:,1], wechsler[:,2], wechsler[:,3:4])
    let out = sprint(show, w)
        @test contains(out, "reject h_0") && !contains(out, "fail to")
    end
    let ci = confint(w)
        @test first(ci) ≈ 0.4963917 atol=1e-6
        @test last(ci) ≈ 0.8447292 atol=1e-6
    end
    @test nobs(w) == 37
    @test dof(w) == 33
    @test pvalue(w) < 0.00001

    X = [ 2 1 0
          4 2 0
         15 3 1
         20 4 1]
    @test @inferred(partialcor(X[:,1], X[:,2], X[:,3])) ≈ 0.919145 atol=1e-6

    x = PartialCorTest(view(X,:,1), view(X,:,2), view(X,:,3))
    @test contains(sprint(show, x), "fail to reject")
    @test confint(x) == (-1.0, 1.0)
    @test nobs(x) == 4
    @test dof(x) == 1
    @test pvalue(x) ≈ 0.25776212 atol=1e-6
end

@testset "One sample Hotelling's T²" begin
    t = OneSampleHotellingT2(nutrient, [1000, 15, 60, 800, 75])
    @test nobs(t) == 737
    @test dof(t) == (5, 732)
    @test pvalue(t) ≈ 0.0 atol=eps()
    @test t.T² ≈ 1758.5413137 atol=1e-6
    @test t.F ≈ 349.7968048 atol=1e-6
    let out = sprint(show, t)
        @test contains(out, "reject h_0") && !contains(out, "fail to")
    end

    # Paired
    p = OneSampleHotellingT2(spouse[:,1:4], spouse[:,5:end])
    @test nobs(p) == 30
    @test dof(p) == (4, 26)
    @test pvalue(p) ≈ 0.039369144 atol=1e-6
    @test p.T² ≈ 13.127840261 atol=1e-6
    @test p.F ≈ 2.942446955 atol=1e-6
    let out = sprint(show, p)
        @test contains(out, "reject h_0") && !contains(out, "fail to")
    end
end

@testset "Two sample Hotelling's T²" begin
    eq = EqualCovHotellingT2(genuine, counterfeit)
    @test nobs(eq) == (100, 100)
    @test dof(eq) == (6, 193)
    @test pvalue(eq) ≈ 0.0 atol=eps()
    @test eq.T² ≈ 2412.4506855 atol=1e-6
    @test eq.F ≈ 391.9217023 atol=1e-6
    @test contains(sprint(show, eq), "reject h_0")

    un = UnequalCovHotellingT2(genuine, counterfeit)
    @test nobs(un) == (100, 100)
    @test dof(un) == (6, 193)
    @test pvalue(un) ≈ 0.0 atol=eps()
    @test un.T² ≈ 2412.4506855 atol=1e-6
    @test un.F ≈ 391.9217023 atol=1e-6
    @test contains(sprint(show, un), "reject h_0")
end

@testset "Bartlett's test" begin
    b = BartlettsTest(genuine, counterfeit)
    @test nobs(b) == (100, 100)
    @test dof(b) == 21
    @test pvalue(b) ≈ 0.0 atol=1e-10
    @test b.L′ ≈ 121.8991235 atol=1e-6
    @test contains(sprint(show, b), "reject h_0")
end


@syms a b c

using Metatheory

@testset "Equality" begin
    @eqtest a == a
    @eqtest a != b
    @eqtest a*b == a*b
    @eqtest a*b != a
    @eqtest a != a*b
end

@testset "Literal Matcher" begin
    r = @rule 1 => 4
    @test r(1) === 4
    @test r(1.0) === 4
    @test r(2) === nothing
end

@testset "Slot matcher" begin
    @test @rule(~x => true)("?") === true
    @test @rule( ~x => ~x)(2) === 2
end

@testset "Term matcher" begin
    @test @rule(sin(~x) => ~x)(sin(a)) === a
    @eqtest @rule(sin(~x) => ~x)(sin(a^2)) == a^2
    @test @rule(sin(~x) => ~x)(sin(a)^2) === nothing
    @test @rule(sin(sin(~x)) => ~x)(sin(a^2)) === nothing
    @test @rule(sin(sin(~x)) => ~x)(sin(sin(a))) === a
    @test @rule(sin(~x)^2 => ~x)(sin(a)^2) === a
end

@testset "Equality matching" begin
    @test @rule((~x)^(~x) => ~x)(a^a) === a
    @test @rule((~x)^(~x) => ~x)(b^a) === nothing
    @test @rule((~x)^(~x) => ~x)(a+a) === nothing
    @eqtest @rule((~x)^(~x) => ~x)(sin(a)^sin(a)) == sin(a)
    @eqtest @rule((~x*~y + ~x*~z)  => ~x * (~y+~z))(a*b + a*c) == a*(b+c)

    @eqtest @rule(+(~~x) => ~~x)(a + b) == [a,b]
    @eqtest @rule(+(~~x) => ~~x)(term(+, a, b, c)) == [a,b,c]
    @eqtest @rule(+(~~x,~y, ~~x) => (~~x, ~y))(term(+,9,8,9,type=Any)) == ([9,],8)
    @eqtest @rule(+(~~x,~y, ~~x) => (~~x, ~y, ~~x))(term(+,9,8,9,9,8,type=Any)) == ([9,8], 9, [9,8])
    @eqtest @rule(+(~~x,~y,~~x) => (~~x, ~y, ~~x))(term(+,6,type=Any)) == ([], 6, [])
end

using SymbolicUtils: @capture

@testset "Capture form" begin

    ex = a^a

    #note that @test inserts a soft local scope (try-catch) that would gobble
    #the matches from assignment statements in @capture macro, so we call it
    #outside the test macro
    ret = @capture ex (~x)^(~x)
    @test ret
    @test @isdefined x
    @test x === a

    ex = b^a
    ret = @capture ex (~y)^(~y)
    @test !ret
    @test !(@isdefined y)

    ret = @capture (a + b) (+)(~~z)
    @test ret
    @test @isdefined z
    @test all(z .=== arguments(a + b))

    #a more typical way to use the @capture macro

    f(x) = if @capture x (~w)^(~w)
        w
    end

    @eqtest f(b^b) == b
    @test f(b+b) == nothing
end


"""
    WeightedVariance(T = Float64)

Simple weighted variance, tracked as type `T`.

# Example:
    o = fit!(WeightedVariance(), rand(100), rand(100))
    sum(o)
    mean(o)
    var(o)
    std(o)
"""
mutable struct WeightedVariance{T} <: WeightedOnlineStat{T}
    μ::T
    σ2::T
    W::T
    W2::T
    n::Int
    function WeightedVariance{T}(
            μ = T(0), σ2 = T(0), W = T(0), W2 = T(0), n = 0
        ) where T
        new{T}(T(μ), T(σ2), T(W), T(W2), Int(n))
    end
end

WeightedVariance(μ::T, σ2::T, W::T, W2::T, n::Int) where T =
    WeightedVariance{T}(μ, σ2, W, W2, n)
WeightedVariance(::Type{T}) where T = WeightedVariance(T(0), T(0), T(0), T(0), 0)
WeightedVariance() = WeightedVariance(Float64)

function OnlineStatsBase._fit!(o::WeightedVariance{T}, x, w) where T
    xx = convert(T, x)
    ww = convert(T, w)

    o.n += 1
    γ1 = T(1) / o.n
    o.W = smooth(o.W, ww, γ1)
    o.W2 = smooth(o.W2, ww * ww, γ1)
    γ = ww / (o.W * o.n)
    μ = o.μ

    o.μ = smooth(o.μ, xx, γ)
    o.σ2 = smooth(o.σ2, (xx - o.μ) * (xx - μ), γ)

    return o
end

function OnlineStatsBase._merge!(
        o::WeightedVariance{T},
        o2::WeightedVariance
    ) where T

    o2_μ = convert(T, o2.μ)
    o2_σ2 = convert(T, o2.σ2)
    o2_W = convert(T, o2.W)
    o2_W2 = convert(T, o2.W2)

    n = o.n + o2.n
    W = smooth(o.W, o2_W, o2.n / n)
    γ1 = (o.W * o.n) / (W * n)
    γ2 = (o2_W * o2.n) / (W * n)

    μ = smooth(o.μ, o2_μ, γ2)

    # o.σ2 =
    #     γ1 * ( o.σ2 +  o.μ ^ 2) +
    #     γ2 * (o2_σ2 + o2_μ ^ 2) -
    #     μ ^ 2
    o.σ2 =
        γ1 * ( o.σ2  + (o.μ - μ) ^ 2) +
        γ2 * (o2_σ2 + (o2_μ - μ) ^ 2)

    o.n = n
    o.μ = μ
    o.W = W
    o.W2 = OnlineStats.smooth(o.W2, o2_W2, o2.n / o.n)

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

    # μ = o.μ
    #
    # γ = o2_W / (o2_W + o.W)
    # δ = o2_μ - o.μ
    #
    # o.σ2 = OnlineStats.smooth(o.σ2, o2_σ2, γ) + (δ ^ 2) * γ * (1.0 - γ)
    # o.μ = OnlineStats.smooth(o.μ, o2_μ, γ)
    # # o.σ2 = o.σ2 + (o.W + o2_W) * (μ)
    #
    # o.W += o2_W
    # o.W2 += o2_W2

    return o
end

OnlineStatsBase.value(o::WeightedVariance) = o.σ2
Base.sum(o::WeightedVariance) = mean(o) * meanweight(o) * nobs(o)
Statistics.mean(o::WeightedVariance) = o.μ
function Statistics.var(
        o::WeightedVariance;
        corrected = true,
        weight_type = :analytic
    )
    if corrected
        if weight_type == :analytic
            value(o) / (1 - (o.W2 * nobs(o)) / (weightsum(o) ^ 2))
        elseif weight_type == :frequency
            value(o) / (weightsum(o) - 1) * weightsum(o)
        elseif weight_type == :probability
            error("If you need this, please make a PR or open an issue")
        else
            throw(ArgumentError("weight type $weight_type not implemented"))
        end
    else
        value(o)
    end
end
Statistics.std(o::WeightedVariance; kw...) = sqrt.(var(o; kw...))
Base.copy(o::WeightedVariance) = WeightedVariance(o.μ, o.σ2, o.W, o.W2, o.n)


default_lambda{T <: Float}(x::DenseMatrix{T}, y::DenseVector{T}) = sqrt( log(size(x,2)) / length(y)) :: T
default_lambda{T <: Float}(x::DenseMatrix{T}, y::DenseVector{T}, q::Int) = (ones(T, q) * sqrt( log(size(x,2)) / length(y))) :: Vector{T}
default_lambda(T::Type, n::Int, p::Int, q::Int) = (ones(T, q)) * sqrt( log(p / n)) :: Vector{T}
default_lambda(n::Int, p::Int, q::Int) = default_lambda(Float64, n, p, q)

"""
    fit_logistic(x, y, λ)

Refit a regularized logistic model using Newton's method.

Arguments:

- `x` is the `n` x `p` statistical model.
- `y` is the `n`-vector of binary responses.
- `λ` is the regularization parameter

Optional Arguments:

- `tol` is the convergence tolerance. Defaults to `1e-8`.
- `max_iter` sets the maximum number of backtracking steps. Defaults to `50`.
- `quiet` controls output. Defaults to `true` (no output).

Output:

- a vector of refit coefficients for `β`
- the number of backtracking steps taken.
"""
function fit_logistic!{T <: Float, V <: DenseVector}(
    v        :: IHTLogVariables{T, V},
    y        :: V,
    λ        :: T;
    tol      :: T    = convert(T, 1e-8),
    max_iter :: Int  = 50,
    quiet    :: Bool = true
)
    # number of cases?
    n = length(y)

    # if b is not warm-started, then ensure that it is not entirely zero
    if all(v.bk .== 0)
        v.bk[1] = logit(abs(mean(y)))
    end

    # initialize intermediate arrays for calculations
    #BLAS.gemv!('N', one(T), x, b, zero(T), xβ)
    A_mul_B!(v.xb, v.xk, v.bk)
    log2xb!(v.lxb, v.l2xb, v.xb)
    copy!(v.bk0, v.bk)
    fill!(v.db, zero(T))

    # track objective
    old_obj = convert(T, Inf)
    # new_obj = (sum(log(1 + exp(v.xb))) - dot(y, v.xb)) / n + λ*sumabs2(v.bk) / 2
    new_obj = logistic_loglik(v.xb, y, v.bk, λ)

    # output progress to console
    quiet || println("Iter\tHalves\tObjective")

    i = 0
    bktrk = 0
    # enter loop for Newton's method
    for i = 1:max_iter

        # db = (x'*(lxβ - y)) / n + λ*b
        BLAS.axpy!(-one(T), sdata(y), sdata(v.lxb))
        At_mul_B!(v.db, v.xk, v.lxb)
        scale!(v.db, 1/n)
        BLAS.axpy!(λ, v.bk, v.db)

        # d2b = (x'*diagm(l2xb)*x)/n + λ*I
        # note that log2xb!() already performs division by n on l2xb
        copy!(v.xk2, v.xk)
        scale!(v.l2xb, v.xk2)
        At_mul_B!(v.d2b, v.xk, v.xk2)
        v.d2b += λ*I

        # b = b0 - ntb = b0 - inv(d2b)*db
        #   = b0 - inv[ x' diagm(pi) diagm(1 - pi) x + λ*I] [x' (pi - y) + λ*b]
        try
            v.ntb = v.d2b \ v.db
        catch e
            warn("in fit_logistic, aborting after caught error: ", e)
            return i, div(bktrk, max_iter)
        end
        copy!(v.bk, v.bk0)
        BLAS.axpy!(-one(T), v.ntb, v.bk)

        # compute objective
        A_mul_B!(v.xb, v.xk, v.bk)
        # new_obj = (sum(log(1 + exp(v.xb))) - dot(y, v.xb)) / n + λ*sumabs2(v.bk) / 2
        new_obj = logistic_loglik(v.xb, y, v.bk, λ)

        # control against nonfinite objective
        isfinite(new_obj) || throw(error("in fit_logistic, objective is no longer finite"))

        # backtrack
        j = 0
        while (new_obj > old_obj + tol) && (j < 50)

            # increment iterator
            j += 1

            # b = b0 - 0.5*ntb
            copy!(v.bk, v.bk0)
            #BLAS.axpy!(p,-one(T) / (2^j),ntb,1,b,1)
            #BLAS.axpy!(-one(T) / (2^j),ntb,v.b)
            BLAS.axpy!(-1/(2^j), v.ntb, v.bk)

            # recalculate objective
            A_mul_B!(v.xb, v.xk, v.bk)
            new_obj = logistic_loglik(v.xb, y, v.bk, λ)

            # ensure a finite objective
            isfinite(new_obj) || throw(error("in fit_logistic, objective is no longer finite"))
        end

        # accumulate total backtracking steps
        bktrk += j

        # track distance between iterates
        dist = euclidean(v.bk, v.bk0) / (norm(v.bk0,2) + one(T))

        # track progress
        quiet || println(i, "\t", j, "\t", dist)

        # check for convergence
        # if converged, then return b
        dist < tol && return i, div(bktrk, i)

        # unconverged at this point, so update intermediate arrays
        #BLAS.gemv!('N', one(T), sdata(x), sdata(b), zero(T), sdata(xβ))
        A_mul_B!(v.xb, v.xk, v.bk)
        log2xb!(v.lxb, v.l2xb, v.xb)

        # save previous β
        copy!(v.bk0, v.bk)
        old_obj = new_obj
    end

    warn("fit_logistic failed to converge in $(max_iter) iterations, exiting...")
    return i, div(bktrk, max_iter)
end

"""
L0 PENALIZED LOGISTIC REGRESSION

    L0_log(x,y,k) -> IHTLogResults

This routine minimizes the loss function given by the negative logistic loglikelihood

    L(β) = sum(log(1 + exp(x*β))) - y'*x*β

subject to `β` lying in the set S_k = { x in R^p : || x ||_0 <= k }.
To ensure a stable model selection process, the optimization is performed over a Tikhonov-regularized copy of `L(β)`; the actual optimized objective is

    g(β) = L(β) + 0.5*lambda*sumabs2(β)

where `lambda` controls the strength of the L2 penalty.
This function extends the [MATLAB source code](http://users.ece.gatech.edu/sbahmani7/GraSP.html) for Sohail Bahmani's nonlinear hard thresholding pursuit framework [GraSP](http://jmlr.csail.mit.edu/papers/v14/bahmani13a.html).

Arguments:

- `x` is the `n` x `p` data matrix.
- `y` is the `n`-dimensional continuous response vector.
- `k` is the desired model size (support).

Optional Arguments:

- `lambda` is the strength of the regularization parameter. Defaults to `sqrt(log(p)/n)`.
- `mu` is the step size used in gradient descent. Defaults to `1.0`.
- `max_iter` is the maximum number of iterations for the algorithm. Defaults to `100`.
- `max_step` is the maximum number of backtracking steps for the step size calculation. Defaults to `100`.
- `tol` is the global tolerance. Defaults to `1e-6`.
- `tolG` is the tolerance for the gradient. Convergence occurs if the norm of the largest `3*k` elements of the gradient falls below `tolG`. Defaults to `1e-6`.
- `tolrefit` is the tolerance for Newton's method in the refitting routine. The refitting algorithm converges when the loss function calculated on the active set (that is, the refit coefficients) falls below `tolrefit`. Defaults to `1e-6`.
- `refit` is a `Bool` that controls refitting at every iterations. It is wise to refit the nonzero coefficients of `b` at every iteration as this may improve convergence behavior and estimation. Defaults to `true` (refit at every iteration).
- `quiet` is a `Bool` that controls algorithm output. Defaults to `true` (no output).

Outputs are wrapped into an `IHTLogResults` object with the following fields:

- 'time' is the compute time for the algorithm. Note that this does not account for time spent initializing optional argument defaults.
- 'iter' is the number of iterations that the algorithm took before converging.
- 'loss' is the optimal loss (half of residual sum of squares) at convergence.
- 'beta' is the final estimate of `b`.
- `active` is the final active set. The size of this set ranges from `k` to `2k`.
"""
function L0_log{T <: Float, V <: DenseVector}(
    x        :: DenseMatrix{T},
    y        :: V,
    k        :: Int;
    v        :: IHTLogVariables{T, V} = IHTLogVariables(sdata(x), sdata(y), k),
    lambda   :: T    = default_lambda(x, y),
    mu       :: T    = one(T),
    tol      :: T    = convert(T, 1e-6),
    tolG     :: T    = convert(T, 1e-3),
    tolrefit :: T    = convert(T, 1e-6),
    max_iter :: Int  = 100,
    max_step :: Int  = 100,
    refit    :: Bool = true,
    quiet    :: Bool = true
)

    # start timer
    tic()

    # problem dimensions?
    n,p = size(x)

    # check arguments
    length(y) == n      || throw(ArgumentError("Length n = $n of response vector y does not match number of rows = $(size(x,1)) in x"))
    k        >  p       && throw(ArgumentError("Value of argument k = $k exceeds number of predictors p = $p"))
    lambda   <  zero(T) && throw(ArgumentError("Value of argument lambda = $lambda must be nonnegative"))
    mu       >  zero(T) || throw(ArgumentError("Value of argument mu = $mu must be positive"))
    tol      >  eps(T)  || throw(ArgumentError("Value of argument tol must exceed machine precision"))
    tolG     >  eps(T)  || throw(ArgumentError("Value of argument tolG must exceed machine precision"))
    tolrefit >  eps(T)  || throw(ArgumentError("Value of argument tolrefit must exceed machine precision"))
    max_iter >= 0       || throw(ArgumentError("Value of max_iter must be nonnegative\n"))
    max_step >= 0       || throw(ArgumentError("Value of max_step must be nonnegative\n"))

    # initialize return values
    iter      = 0                 # number of iterations of L0_reg
    exec_time = zero(T)           # compute time *within* L0_reg
    loss      = convert(T, Inf)   # loss function value

    # initialize algorithm parameters
    num_df    = min(p, 3*k)       # largest 3*k active components of gradient
    short_df  = min(p, 2*k)       # largest 2*k active components of gradient
    converged = false             # is algorithm converged?
    stuck     = false             # is algorithm stuck in a cycle?
    normdf    = convert(T, Inf)   # norm of active portion of gradient
    loss0     = convert(T, Inf)   # penultimate loss function value
    loss00    = convert(T, Inf)   # antepenultimate loss function value
    bktrk     = 0                 # number of Newton backtracking steps
    nt_iter   = 0                 # number of Newton iterations
    lt        = length(v.active)  # size of current active set?

    # formatted output to monitor algorithm progress
    quiet || print_header_log()

    # main GraSP iterations
    for iter = 1:max_iter

        # notify and break if maximum iterations are reached
        # also break if algorithm is cycling
        if iter >= max_iter || stuck

            # warn about hitting maximum iterations
            iter >= max_iter && print_maxiter(max_iter, loss)

            # warn about cycling
            stuck && warn("L0_log appears to be cycling after $iter iterations, aborting early...\nCurrent loss: $loss\n")

            # send elements below tol to zero
            threshold!(v.b, tol)

            # if requested, apply final refit without regularization
            if refit
                copy!(v.idxs0, v.idxs)
                update_indices!(v.idxs, v.b)
                !isequal(v.idxs, v.idxs0) && copy!(v.xk, view(x, :, v.idxs))
                nt_iter, bktrk = fit_logistic!(v, y, zero(T), tol=tolrefit, max_iter=max_step, quiet=true)
                v.b[v.idxs] = v.bk
            end

            # stop timer
            exec_time = toq()

            # these are output variables for function
            # wrap them into a Dict and return
            return IHTLogResults(exec_time, iter, loss, copy(v.b), copy(v.active))
        end

        # save previous loss, iterate
        loss00 = loss0
        loss0  = loss
        copy!(v.b0, v.b)

        # size of current active set?
        if iter > 1
            v.active = union(v.dfidxs[1:short_df], v.bidxs[1:k])
        end
        lt = length(v.active)

        # update x*b
        # no need to compute anything if b = 0
        if all(v.b .== 0)
            fill!(v.xb, zero(T))
        else
            # A_mul_B!(v.xb, x, v.b)
            update_xb!(v.xb, x, v.b, v.active, lt)
        end

        # recompute active loss = (-dot(y,Xb) + sum(log(1.0 + exp(Xb)))) / n + 0.5*lambda*sumabs2(b[active])
        # special case: b = 0 --> Xb = 0 --> loss = n*log(1 + exp(0))/n + 0.5*lambda*norm(0)
        if iter < 2
            loss = 1 / eps(T)
        elseif all(v.xb .== 0)
            loss = log(2)
        else
            loss = logistic_loglik(v.xb, y, v.b, v.active, lambda, k)
        end

        # guard against numerical instabilities in loss function
        check_finiteness(loss)

        # recompute active gradient df[active] = (x[:,active]'*(logistic(Xb) - y)) / n + lambda*b[active]
        # arrange calculations differently if active set is entire support 1, 2, ..., p
        if lt == p
            logistic_grad!(v.df, v.lxb, x, y, v.b, v.xb, lambda)
        else
            logistic_grad!(v.df, v.lxb, x, y, v.b, v.xb, v.active, lt, lambda)
        end
#        logistic!(v.lxb, v.xb)
#        v.df = (x' * (v.lxb - y)) / n + lambda*v.b

        # need norm of top 3*k components of gradient
        normdf = df_norm(v.df, v.dfidxs, 1, num_df)

        # guard against numerical instabilities in gradient
        check_finiteness(normdf)

        # identify 2*k dominant directions in gradient
        selectperm!(v.dfidxs, v.df, 1:num_df, by=abs, rev=true, initialized=true)

        # clean b and fill, b[active] = b0[active] - μ*df[active]
        # note that sparsity level is size(active) which is one of [k, k+1, ..., 3*k]
        #update_x!(v.b, v.b0, v.df, v.active, mu)
        fill!(v.b, zero(T))
        v.b[v.active] = v.b0[v.active] .- mu.*v.df[v.active]

        # now apply hard threshold on model to original desired sparsity k
        project_k!(v.b, k)
        selectperm!(v.bidxs, v.b, 1:k, by=abs, rev=true, initialized=true)
        v.bk = v.b[view(v.bidxs, 1:k)]
        fill!(v.b, 0)
        v.b[view(v.bidxs, 1:k)] = v.bk
        update_xb!(v.xb, x, v.b, v.active, lt)
        loss = logistic_loglik(v.xb, y, v.b, v.active, lambda, k)

        # check for convergence
        converged_obj  = abs(loss - loss0) < tol
        converged_grad = normdf < tolG
        converged      = converged_obj || converged_grad
        stuck          = abs(loss - loss00) < tol
#        stuck          = abs(loss - loss00) < tol && converged_grad

        # output algorithm progress
        quiet || @printf("%d\t%d\t%d\t%3.7f\t%3.7f\n",iter,nt_iter,bktrk,loss,normdf)

        # check for convergence
        # if converged and in feasible set, then algorithm converged before maximum iteration
        # perform final computations and output return variables
        if converged

            # send elements below tol to zero
            threshold!(v.b, tol)

            # if requested, apply final refit without regularization
            # note that regularization is NOT set exactly to zero!
            # this aids numerical stability
            if refit
                copy!(v.idxs0, v.idxs)
                update_indices!(v.idxs, v.b)
#                !isequal(idxs,idxs0) && update_xk!(xk, x, idxs, k=k, n=n, p=p)
                copy!(v.xk, view(x, :, v.idxs))
                try
                    nt_iter, bktrk = fit_logistic!(v, y, 2*eps(T), tol=tolrefit, max_iter=max_step, quiet=true)
                    v.b[v.idxs] = v.bk
                catch e
#                    warn("in final refitting, caught error: ", e)
#                    warn("skipping final refit")
                    nt_iter = 0
                    bktrk   = 0
                end # end try-catch for refit
            end

            # stop time
            exec_time = toq()

            # announce convergence
            !quiet && print_log_convergence(iter, loss, exec_time, normdf)

            # return output
            return IHTLogResults(exec_time, iter, loss, copy(v.b), copy(v.active))
        end # end convergence check

        # unconverged at this point
        # refit nonzeroes in b?
        if refit

            # update boolean vector of nonzeroes
            copy!(v.idxs0, v.idxs)
            update_indices!(v.idxs, v.b)

            # update active set of x, if necessary
#            (iter == 1 || !isequal(idxs,idxs0)) && update_xk!(xk, x, idxs, k=k, n=n, p=p)
            copy!(v.xk, view(x, :, v.idxs))
            copy!(v.xk2, v.xk)

            # attempt refit but guard against possible instabilities or singularities
            # these are more frequent if lambda is small
            # if refitting destabilizes, then leave b alone
            try
                nt_iter, bktrk = fit_logistic!(v, y, lambda, tol=tolrefit, max_iter=max_step, quiet=true)
                v.b[v.idxs] = v.bk
            catch e
#                warn("in refitting, caught error: ", e)
#                warn("skipping refit")
                nt_iter = 0
                bktrk   = 0
            end # end try-catch for refit
        end # end refit

    end # end main GraSP iterations

    # return null result
    return IHTLogResults(zero(T), 0, zero(T), zeros(T, 1), zeros(Int, 1))
end # end L0_log

"""
COMPUTE AN IHT REGULARIZATION PATH FOR LOGISTIC REGRESSION

    iht_path_log(x,y,path) -> SparseCSCMatrix

This subroutine computes best logistic models for matrix `x` and response `y` by calling `L0_log` for each model over a regularization path denoted by `path`.

Arguments:

- `x` is the `n` x `p` design matrix.
- `y` is the `n`-vector of responses.
- `path` is an `Int` vector that contains the model sizes to test.

Optional Arguments:

- `tol` is the global convergence tolerance for `L0_log`. Defaults to `1e-6`.
- `tolG` is the tolerance for the gradient. Convergence occurs if the norm of the largest `3*k` elements of the gradient falls below `tolG`. Defaults to `1e-6`.
- `tolrefit` is the tolerance for Newton's method in the refitting routine. The refitting algorithm converges when the loss function calculated on the active set (that is, the refit coefficients) falls below `tolrefit`. Defaults to `1e-6`.
- `max_iter` caps the number of iterations for the algorithm. Defaults to `100`.
- `max_step` caps the number of backtracking steps in the IHT kernel. Defaults to `100`.
- `refit` is a `Bool` that controls refitting at every iterations. It is wise to refit the nonzero coefficients of `b` at every iteration as this may improve convergence behavior and estimation. Defaults to `true` (refit at every iteration).
- `quiet` is a Boolean that controls the output. Defaults to `true` (no output).

Output:

- a sparse `p` x `length(path)` matrix where each column contains the computed model for each component of `path`.
"""
function iht_path_log{T <: Float}(
    x        :: DenseMatrix{T},
    y        :: DenseVector{T},
    path     :: DenseVector{Int};
    lambdas  :: Vector{T} = default_lambda(x, y, length(path)),
    mu       :: T    = one(T),
    tol      :: T    = convert(T, 1e-6),
    tolG     :: T    = convert(T, 1e-3),
    tolrefit :: T    = convert(T, 1e-6),
    max_iter :: Int  = 100,
    max_step :: Int  = 50,
    refit    :: Bool = true,
    quiet    :: Bool = true,
)

    # dimensions of problems?
    n,p = size(x)

    # how many models will we compute?
    num_models = length(path)

    # preallocate temporary arrays
    v = IHTLogVariables(sdata(x), sdata(y), 1)

    # betas will be a sparse matrix with models
    betas = spzeros(T, p, num_models)

    # compute the path
    for i = 1:num_models

        # model size?
        q = path[i]
        quiet || println("Current model size: $q")

        # update temporary variables with new threshold size q
        update_variables!(v, x, q)

        # store projection of beta onto largest k nonzeroes in magnitude
        project_k!(v.b, q)

        # current regularization parameter?
        λ = lambdas[i]

        # now compute current model
        output = L0_log(x, y, q, v=v, lambda=λ, mu=mu, tol=tol, tolG=tolG, tolrefit=tolrefit, max_iter=max_iter, max_step=max_step, refit=refit, quiet=quiet)

        # extract and save model
        betas[:,i] = sparsevec(output.beta)

        # run garbage collector to clear temporary arrays for next step of path
        #gc()
    end

    # return a sparsified copy of the models
    return betas
end

"""
COMPUTE ONE FOLD IN A CROSSVALIDATION SCHEME FOR A PENALIZED LOGISTIC REGRESSION REGULARIZATION PATH

    one_fold_log(x,y,path,folds,fold) -> Vector{Float}

For a regularization path given by the `Int` vector `path`,
this function performs penalized logistic regression on `x` and `y` and computes an out-of-sample error based on the indices given in `folds`.

Arguments:

- `x` is the `n` x `p` design matrix.
- `y` is the `n`-vector of responses.
- `path` is an `Int` vector that specifies which model sizes to include in the path, e.g. `path = collect(k0:increment:k_end)`.
- `folds` is an `Int` vector indicating which component of `y` goes to which fold, e.g. `folds = IHT.cv_get_folds(n,nfolds)`
- `fold` is the current fold to compute.

Optional Arguments:

- `tol` is the convergence tolerance to pass to the path computations. Defaults to `1e-4`.
- `max_iter` caps the number of permissible iterations in the IHT algorithm. Defaults to `100`.
- `max_step` caps the number of permissible backtracking steps. Defaults to `100`.
- `quiet` is a Boolean to activate output. Defaults to `true` (no output).
- `deviance` is an `String` that specifies the crossvalidation criterion, either `"deviance"` for the logistic loglikelihood or `"class"` for misclassification error. Defaults to `"deviance"`.

Output:

- `errors` is a vector of out-of-sample errors (MSEs) for the current fold.
"""
function one_fold_log{T <: Float}(
    x         :: DenseMatrix{T},
    y         :: DenseVector{T},
    path      :: DenseVector{Int},
    folds     :: DenseVector{Int},
    fold      :: Int;
    lambdas   :: Vector{T} = default_lambda(x, y, length(path)),
    criterion :: String = "deviance",
    tol       :: T      = convert(T, 1e-6),
    tolG      :: T      = convert(T, 1e-3),
    tolrefit  :: T      = convert(T, 1e-6),
    max_iter  :: Int    = 100,
    max_step  :: Int    = 100,
    refit     :: Bool   = true,
    quiet     :: Bool   = true,
)
    # dimension of problem?
    n,p = size(x)
    nmodels = length(path)

    # ensure that crossvalidation criterion is valid
    criterion in ["deviance", "class"] || throw(ArgumentError("Argument criterion must be 'deviance' or 'class'"))

    # make vector of indices for folds
    # train_idx is the vector that indexes the TRAINING set
    train_idx = folds .!= fold
    test_size = sum(!train_idx)

    # allocate the arrays for the training set
    x_train = x[train_idx,:]
    y_train = y[train_idx]

    # compute the regularization path on the training set
    betas = iht_path_log(x_train, y_train, path, tol=tol, tolG=tolG, tolrefit=tolrefit, max_iter=max_iter, max_step=max_step, refit=refit, quiet=quiet, lambdas=lambdas)

    # (masked) estimated responses
    # mask based on test_idx
    xβ  = x*betas
    fill!(view(xβ, train_idx, :), zero(T))
    lxb = zeros(T, size(xβ))
    logistic!(lxb, xβ)
    fill!(view(lxb, train_idx, :), zero(T))
    ytest  = copy(y)
    fill!(view(ytest, train_idx), zero(T))

    # compute the mean out-of-sample error for the TEST set
    errors = zeros(T, nmodels)
    if criterion == "class" # compute misclassification error
        errors .= vec(sumabs(lxb .- ytest, 1)) ./ sum(!train_idx)
    else # else -> criterion == "deviance" = 2*(neg logistic loss with no reg)
        #errors .= - one(T) ./ n .* ( vec( sum(xβ .* y, 1) .- sum(log(one(T) .+ exp(xβ)), 1) ) )
        errors .= - 2 ./ test_size .* ( vec( sum(xβ .* ytest, 1) .- sum(log(one(T) .+ exp(xβ)), 1) ) )
    end
    return errors :: Vector{T}
end

#"""
#    pfold_log(xfile, xtfile, x2file, yfile, meanfile, invstdfile, path, folds, numfolds[, pids=procs()])
#
#This function is the parallel execution kernel in `cv_log()`. It is not meant to be called outside of `cv_iht()`.
#It will distribute `numfolds` crossvalidation folds across the processes supplied by the optional argument `pids` and call `one_fold_log()` for each fold.
#Each fold will compute a regularization path given by `path`.
#`pfold()` collects the vectors of MSEs returned by calling `one_fold_log()` for each process, reduces them, and returns their average across all folds.
#"""
function pfold_log{T <: Float}(
    x         :: DenseMatrix{T},
    y         :: DenseVector{T},
    path      :: DenseVector{Int},
    folds     :: DenseVector{Int},
    q         :: Int;
    pids      :: Vector{Int} = procs(),
    lambdas   :: DenseVector{T} = default_lambda(x, y, length(path)),
    criterion :: String = "deviance",
    tol       :: T      = convert(T, 1e-6),
    tolG      :: T      = convert(T, 1e-3),
    tolrefit  :: T      = convert(T, 1e-6),
    max_iter  :: Int    = 100,
    max_step  :: Int    = 100,
    quiet     :: Bool   = true,
    refit     :: Bool   = true,
    header    :: Bool   = false
)

    # problem size?
    n,p = size(x)

    # ensure that crossvalidation criterion is valid
    criterion in ["deviance", "class"] || throw(ArgumentError("Argument criterion must be 'deviance' or 'class'"))

    # how many CPU processes can pfold use?
    np = length(pids)

    # report on CPU processes
    quiet || println("pfold: np = ", np)
    quiet || println("pids = ", pids)

    # set up function to share state (indices of folds)
    i = 1
    nextidx() = (idx=i; i+=1; idx)

    # preallocate cell array for results
    results = SharedArray(T, (length(path),q), pids=pids) :: SharedMatrix{T}

    # master process will distribute tasks to workers
    # master synchronizes results at end before returning
    @sync begin

        # loop over all workers
        for worker in pids

            # exclude process that launched pfold, unless only one process is available
            if worker != myid() || np == 1

                # asynchronously distribute tasks
                @async begin
                    while true

                        # grab next fold
                        current_fold = nextidx()

                        # if current fold exceeds total number of folds then exit loop
                        current_fold > q && break

                        # report distribution of fold to worker and device
                        quiet || print_with_color(:blue, "Computing fold $current_fold on worker $worker.\n\n")

                        # launch job on worker
                        # worker loads data from file paths and then computes the errors in one fold
                        r = remotecall_fetch(worker) do
                                one_fold_log(x, y, path, folds, current_fold, max_iter=max_iter, max_step=max_step, quiet=quiet, tol=tol, tolG=tolG, tolrefit=tolrefit, refit=refit, criterion=criterion, lambdas=lambdas)
                        end # end remotecall_fetch()
                        setindex!(results, r, :, current_fold)
                    end # end while
                end # end @async
            end # end if
        end # end for
    end # end @sync

    # return reduction (row-wise sum) over results
    return (vec(sum(results, 2) ./ q)) :: Vector{T}
end

"""
PARALLEL CROSSVALIDATION ROUTINE FOR GRASP ALGORITHM FOR PENALIZED LOGISTIC REGRESSION

    cv_iht_log(x,y,path,q) -> Vector{Float}

This function will perform `q`-fold cross validation for the ideal model size in GraSP logistic regression.
It computes several paths as specified in the `path` argument using the design matrix `x` and the response vector `y`.
Each path is asynchronously spawned using any available processor.
For each path, one fold is held out of the analysis for testing, while the rest of the data are used for training.
The function to compute each path, `one_fold_log()`, will return a vector of out-of-sample errors (MCEs).

Arguments:

- `x` is the `n` x `p` design matrix.
- `y` is the `n`-vector of responses.
- `path` is an `Int` vector that specifies which model sizes to include in the path, e.g. `path = collect(k0:increment:k_end)`.
- `q` is the number of folds to compute.

Optional Arguments:

- `n` is the number of samples. Defaults to `length(y)`.
- `p` is the number of predictors. Defaults to `size(x,2)`.
- `folds` is the partition of the data. Defaults to `IHT.cv_get_folds(n,q)`.
- `tol` is the convergence tolerance to pass to the path computations. Defaults to `1e-4`.
- `max_iter` caps the number of permissible iterations in the IHT algorithm. Defaults to `100`.
- `max_step` caps the number of permissible backtracking steps. Defaults to `100`.
- `quiet` is a Boolean to activate output. Defaults to `true` (no output).
   *NOTA BENE*: each processor outputs feed to the console without regard to the others,
   so setting `quiet = false` can yield very messy output!
- `refit` is a `Bool` to indicate whether or not to recompute the best model. Note that this argument affects both the continuous refitting in the GraSP algorithm in addition to a final refit after crossvalidation. Defaults to `true` (recompute).
- `deviance` is an `String` that specifies the crossvalidation criterion, either `"deviance"` for the logistic loglikelihood or `"class"` for misclassification error. Defaults to `"deviance"`.

Output:

- `errors` is the averaged MSE over all folds.

If called with `refit = true`, then the output also includes, for best model size `k_star`:

- `b`, a vector of `k_star` floats
- `bidx`, a vector of `k_star` indices indicating the support of the best model.
"""
function cv_log{T <: Float}(
    x         :: DenseMatrix{T},
    y         :: DenseVector{T};
    q         :: Int              = cv_get_num_folds(3, 5),
    path      :: DenseVector{Int} = collect(1:min(size(x,2),20)),
    folds     :: DenseVector{Int} = cv_get_folds(sdata(y),q),
    pids      :: Vector{Int}      = procs(),
    lambdas   :: Vector{T}        = default_lambda(x, y, length(path)),
    criterion :: String = "deviance",
    tol       :: T      = convert(T, 1e-6),
    tolG      :: T      = convert(T, 1e-3),
    tolrefit  :: T      = convert(T, 1e-6),
    max_iter  :: Int    = 100,
    max_step  :: Int    = 100,
    quiet     :: Bool   = true,
    refit     :: Bool   = true
)

    # problem size?
    n,p = size(x)
    @assert n == length(y)

    # ensure that crossvalidation criterion is valid
    criterion in ["deviance", "class"] || throw(ArgumentError("Argument criterion must be 'deviance' or 'class'"))

    # how many elements are in the path?
    num_models = length(path)

    # compute crossvalidation deviances
    errors = pfold_log(x, y, path, folds, q, max_iter=max_iter, max_step=max_step, quiet=quiet, pids=pids, tolG=tolG, tolrefit=tolrefit, refit=refit, criterion=criterion, lambdas=lambdas)

    # what is the best model size?
    k = path[indmin(errors)] :: Int

    # print results
    !quiet && print_cv_results(errors, path, k)

    # recompute ideal model
    # get λ value for best model
    λ = lambdas[path .== k][1]

    # use L0_log to extract model
    # with refit = true, L0_log will continuously refit predictors
    # no final refitting code necessary
    output = L0_log(x, sdata(y), k, max_iter=max_iter, max_step=max_step, quiet=quiet, tol=tol, tolG=tolG, tolrefit=tolrefit, refit=true, lambda=zero(T))#lambda=λ)

    # which components of beta are nonzero?
    bidx = find(output.beta)

    return IHTCrossvalidationResults(errors, sdata(path), output.beta[bidx], bidx, k)
end


#using DeformableRegistration.ImageProcessing
using DeformableRegistration: Regularizer, Transformation
 using Base.Test

# @testset "check if constant deformation yields zero" begin
#     omega = [0.0,1.0,0.0,1.0];
#     m = [5; 7];
#     h = [1.0,1.0];
#     u = ones(size(getStaggeredGrid(omega,m)))
#     B = createElasticOperatorStaggered(h,m)
#     S, dS, d2S = regularizer(u, B)
#     @test_approx_eq(S[1], 0)
#     @test sum(dS) == 0.0
# end
#
# @testset "check if random deformation yields nonzero result" begin
#     omega = [0.0,1.0,0.0,1.0];
#     m = [5; 7];
#     h = [1.0,1.0];
#     u = rand(size(getStaggeredGrid(omega,m)))
#     B = createElasticOperatorStaggered(h,m)
#     S, dS, d2S = regularizer(u, B)
#     @test S[1] > 0
#     @test (dS'*dS)[1] > 0
# end

@testset "Regularizer" begin
@testset "check if constant deformation yields zero" begin
    m = [5; 7];
    h = [1.0,1.0];
    u = ones(2*m[1],m[2])[:]
    B = createElasticOperatorCentered(h,m)
    S, dS, d2S = regularizer(u, B)
    @test  S[1] ≈ 0
    @test sum(dS) == 0.0
end

@testset "check if random deformation yields nonzero result" begin
    omega = [0.0,1.0,0.0,1.0];
    m = [5; 7];
    h = [1.0,1.0];
    u = rand(2*m[1],m[2])[:]
    B = createElasticOperatorCentered(h,m)
    S, dS, d2S = regularizer(u, B)
    @test S[1] > 0
    @test (dS'*dS)[1] > 0
end

@testset "check if constant deformation yields zero" begin
    omega = [0.0,1.0,0.0,1.0];
    m = [5; 7];
    h = [1.0,1.0];
    u = ones(2*m[1],m[2])[:]
    B = createDiffusiveOperatorCentered(h,m)
    S, dS, d2S = regularizer(u, B)
    @test S[1] ≈ 0
    @test sum(dS) == 0.0
end

@testset "check if constant deformation yields zero" begin
    omega = [0.0,1.0,0.0,1.0];
    m = [5; 7];
    h = [1.0,1.0];
    u = ones(m[1],2*m[2])
    #u[3,3] = 1
    u = u[:]
    B = createCurvatureOperatorCentered(h,m)
    S, dS, d2S = regularizer(u, B)
    @test S[1] ≈ 0
    @test sum(dS) == 0.0

    #using PyPlot
    #imshow(reshape(dS[1:35],(5,7)))
    #plot(dS[1:25])
end

@testset "check if random deformation yields nonzero result" begin
    omega = [0.0,1.0,0.0,1.0];
    m = [5; 7];
    h = [1.0,1.0];
    u = rand(2*m[1],m[2])[:]
    B = createCurvatureOperatorCentered(h,m)
    S, dS, d2S = regularizer(u, B)
    @test S[1] > 0
    @test (dS'*dS)[1] > 0
end

@testset "check if random deformation yields nonzero result" begin
    omega = [0.0,1.0,0.0,1.0];
    m = [5; 7];
    h = [1.0,1.0];
    u = rand(2*m[1],m[2])[:]
    B = createDiffusiveOperatorCentered(h,m)
    S, dS, d2S = regularizer(u, B)
    @test S[1] > 0
    @test (dS'*dS)[1] > 0
end
end #testset regularizer
# speed test
#omega = [0.0,1.0,0.0,1.0];
#m = [512; 512];
#h = [1.0,1.0];
#u = rand(size(getStaggeredGrid(omega,m)))
#@time B = createElasticOperatorStaggered(h,m)
#@time S, dS, d2S = regularizer(u, B)

# visualization
#using PyPlot
#spy(B)


#=
    Decomposition of global constraint matrix_element in Julia ConstraintSolver.jl

    matrix_element(model,x,i,j,val)

    ensures that  x[i,j] = val

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

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

#
# matrix_element is in constraints_utils.jl
#
function matrix_element_test(n=4,print_solutions=true,all_solutions=true)

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

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

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

                                                            # "branch_split"=>:Smallest,
                                                            "branch_split"=>:Biggest, # fastest for this simple model!
                                                            # "branch_split"=>:InHalf,

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

                                                            "time_limit"=>3,

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

    m = n
    @variable(model, 1 <= x[1:n,1:m] <= n, Int)

    # @constraint(model, x[:] in CS.AllDifferent())
    @variable(model, 1 <= ii <= n, Int)
    @variable(model, 1 <= jj <= m, Int)
    @variable(model, 1 <= val <= n, Int)

    for i in 1:n
        @constraint(model, x[i,:] in CS.AllDifferent())
        @constraint(model, x[:,i] in CS.AllDifferent())
    end

    @constraint(model, val == n)
    # @constraint(model, ii == 2)
    # @constraint(model, jj == 3)

    matrix_element(model,x,ii,jj,val)

    # @objective(model,Min,x[n,m])

    # Solve the problem
    optimize!(model)

    status = JuMP.termination_status(model)
    # println("status:$status")
    if status == MOI.OPTIMAL
        num_sols = MOI.get(model, MOI.ResultCount())
        println("num_sols:$num_sols\n")
        if print_solutions
            for sol in 1:num_sols
                println("solution #$sol")
                xx = convert.(Integer,JuMP.value.(x; result=sol))
                iix = convert.(Integer,JuMP.value.(ii; result=sol))
                jjx = convert.(Integer,JuMP.value.(jj; result=sol))
                valx = convert.(Integer,JuMP.value.(val; result=sol))

                # print_grid(xx)
                display(xx)
                println("ii:$iix jj:$jjx val:$valx")
                println()

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

    return status
end

print_solutions=true
all_solutions=false
for n in 2:20
    println("\nn:$n")
    @time status = matrix_element_test(n,print_solutions,all_solutions)
    if status != MOI.OPTIMAL
        break
    end
end


### A Pluto.jl notebook ###
# v0.12.7

using Markdown
using InteractiveUtils

# ╔═╡ 6fd616f6-1ea1-11eb-3814-8bfb4a096c49
using Pkg, DrWatson

# ╔═╡ 6ffe2628-1ea1-11eb-24ea-57f985146a72
begin
	@quickactivate "StatisticsWithJuliaPlutoNotebooks"
	using Plots, LaTeXStrings
end;

# ╔═╡ ed174bc4-1ea0-11eb-1e2f-a32874cec549
md"## Listing 3.6"

# ╔═╡ 70179b62-1ea1-11eb-13ce-7fb4d54bad4b
begin
	xGrid = 0:0.01:10
	uGrid = 0:0.01:1
	busy = 0.8
end

# ╔═╡ 80359808-20f8-11eb-01d8-a96d178126b4
begin
	F(t)= t<=0 ? 0 : 1 - busy*exp(-(1-busy)t)
	infimum(B) = isempty(B) ? Inf : minimum(B)
	invF(u) = infimum(filter((x) -> (F(x) >= u),xGrid))
end

# ╔═╡ 805b0906-20f8-11eb-1142-e99c34f90ba2
begin
	p1 = plot(xGrid, F.(xGrid), c=:blue, xlims=(-0.1,10), ylims=(0,1),
		xlabel=L"x", ylabel=L"F(x)")

	p2 = plot(uGrid, invF.(uGrid), c=:blue, xlims=(0,0.95), ylims=(0,maximum(xGrid)),
		xlabel=L"u", ylabel=L"F^{-1}(u)")

	plot(p1, p2, legend=:none, size=(800, 400))
end

# ╔═╡ 70182cd8-1ea1-11eb-094d-8d0d49cf15f3
md"## End of listing 3.6"

# ╔═╡ Cell order:
# ╟─ed174bc4-1ea0-11eb-1e2f-a32874cec549
# ╠═6fd616f6-1ea1-11eb-3814-8bfb4a096c49
# ╠═6ffe2628-1ea1-11eb-24ea-57f985146a72
# ╠═70179b62-1ea1-11eb-13ce-7fb4d54bad4b
# ╠═80359808-20f8-11eb-01d8-a96d178126b4
# ╠═805b0906-20f8-11eb-1142-e99c34f90ba2
# ╟─70182cd8-1ea1-11eb-094d-8d0d49cf15f3


# This file is part of the TaylorIntegration.jl package; MIT licensed

# jetcoeffs!
@doc doc"""
    jetcoeffs!(eqsdiff::Function, t, x, params)

Returns an updated `x` using the recursion relation of the
derivatives obtained from the differential equations
``\dot{x}=dx/dt=f(x, p, t)``.

`eqsdiff` is the function defining the RHS of the ODE,
`x` contains the Taylor1 expansion of the dependent variable(s) and
`t` is the independent variable, and `params` are the parameters appearing on the
function defining the differential equation. See [`taylorinteg`](@ref) for examples
and convention for `eqsdiff`. Note that `x` is of type `Taylor1{U}` where
`U<:Number`; `t` is of type `Taylor1{T}` where `T<:Real`.

Initially, `x` contains only the 0-th order Taylor coefficient of
the current system state (the initial conditions), and `jetcoeffs!`
computes recursively the high-order derivates back into `x`.

"""
function jetcoeffs!(eqsdiff::Function, t::Taylor1{T}, x::Taylor1{U}, params) where
        {T<:Real, U<:Number}
    order = x.order
    for ord in 0:order-1
        ordnext = ord+1

        # Set `taux`, `xaux`, auxiliary Taylor1 variables to order `ord`
        @inbounds taux = Taylor1( t.coeffs[1:ordnext] )
        @inbounds xaux = Taylor1( x.coeffs[1:ordnext] )

        # Equations of motion
        dx = eqsdiff(xaux, params, taux)

        # Recursion relation
        @inbounds x[ordnext] = dx[ord]/ordnext
    end
    nothing
end

@doc doc"""
    jetcoeffs!(eqsdiff!::Function, t, x, dx, xaux, params)

Mutates `x` in-place using the recursion relation of the
derivatives obtained from the differential equations
``\dot{x}=dx/dt=f(x, p, t)``.

`eqsdiff!` is the function defining the RHS of the ODE,
`x` contains the Taylor1 expansion of the dependent variables and
`t` is the independent variable, and `params` are the parameters appearing on the
function defining the differential equation. See [`taylorinteg`](@ref) for examples
and convention for `eqsdiff`. Note that `x` is of type `Vector{Taylor1{U}}`
where `U<:Number`; `t` is of type `Taylor1{T}` where `T<:Real`. In this case,
two auxiliary containers `dx` and `xaux` (both of the same type as `x`) are
needed to avoid allocations.

Initially, `x` contains only the 0-th order Taylor coefficient of
the current system state (the initial conditions), and `jetcoeffs!`
computes recursively the high-order derivates back into `x`.

"""
function jetcoeffs!(eqsdiff!::Function, t::Taylor1{T},
        x::AbstractArray{Taylor1{U}, N}, dx::AbstractArray{Taylor1{U}, N},
        xaux::AbstractArray{Taylor1{U}, N}, params) where {T<:Real, U<:Number, N}
    order = x[1].order
    for ord in 0:order-1
        ordnext = ord+1

        # Set `taux`, auxiliary Taylor1 variable to order `ord`
        @inbounds taux = Taylor1( t.coeffs[1:ordnext] )
        # Set `xaux`, auxiliary vector of Taylor1 to order `ord`
        for j in eachindex(x)
            @inbounds xaux[j] = Taylor1( x[j].coeffs[1:ordnext] )
        end

        # Equations of motion
        eqsdiff!(dx, xaux, params, taux)

        # Recursion relations
        for j in eachindex(x)
            @inbounds x[j][ordnext] = dx[j][ord]/ordnext
        end
    end
    nothing
end

# __jetcoeffs
"""
    __jetcoeffs!(::Val{bool}, f, t, x, params, tmpTaylor, arrTaylor)
    __jetcoeffs!(::Val{bool}, f, t, x, dx, xaux, params, tmpTaylor, arrTaylor)

Chooses a method of [`jetcoeffs!`](@ref) (hard-coded) or the generated by
[`@taylorize`](@ref)) depending on `Val{bool}` (`bool::Bool`).
"""
@inline __jetcoeffs!(::Val{false}, f, t, x, params, tmpTaylor, arrTaylor) =
    jetcoeffs!(f, t, x, params)
@inline __jetcoeffs!(::Val{true},  f, t, x, params, tmpTaylor, arrTaylor) =
    jetcoeffs!(Val(f), t, x, params, tmpTaylor, arrTaylor)
@inline __jetcoeffs!(::Val{false}, f, t, x, dx, xaux, params, tmpTaylor, arrTaylor) =
    jetcoeffs!(f, t, x, dx, xaux, params)
@inline __jetcoeffs!(::Val{true},  f, t, x, dx, xaux, params, tmpTaylor, arrTaylor) =
    jetcoeffs!(Val(f), t, x, dx, params, tmpTaylor, arrTaylor)

# _allocate_jetcoeffs! fallbacks; needed to define the corresponding parsed method
_allocate_jetcoeffs!(::Taylor1{T}, x::Taylor1{S}, params) where {T,S} =
    (typeof(x)[], Vector{typeof(x)}[])
_allocate_jetcoeffs!(::Taylor1{T}, x::AbstractArray{Taylor1{S}, N},
        ::AbstractArray{Taylor1{S}, N}, params) where {T,S,N} =
    (eltype(x)[], Vector{eltype(x)}[])

# stepsize
"""
    stepsize(x, epsilon) -> h

Returns a maximum time-step for a the Taylor expansion `x`
using a prescribed absolute tolerance `epsilon` and the last two
Taylor coefficients of (each component of) `x`.

Note that `x` is of type `Taylor1{U}` or `Vector{Taylor1{U}}`, including
also the cases `Taylor1{TaylorN{U}}` and `Vector{Taylor1{TaylorN{U}}}`.

Depending of `eltype(x)`, i.e., `U<:Number`, it may be necessary to overload
`stepsize`, specializing it on the type `U`, to avoid type instabilities.
"""
function stepsize(x::Taylor1{U}, epsilon::T) where {T<:Real, U<:Number}
    R = promote_type(typeof(norm(constant_term(x), Inf)), T)
    ord = x.order
    h = convert(R, Inf)
    z = zero(R)
    for k in (ord-1, ord)
        @inbounds aux = norm( x[k], Inf)
        aux == z && continue
        aux1 = _stepsize(aux, epsilon, k)
        h = min(h, aux1)
    end
    return h::R
end

function stepsize(q::AbstractArray{Taylor1{U}, N}, epsilon::T) where
        {T<:Real, U<:Number, N}
    R = promote_type(typeof(norm(constant_term(q[1]), Inf)), T)
    h = convert(R, Inf)
    for i in eachindex(q)
        @inbounds hi = stepsize( q[i], epsilon )
        h = min( h, hi )
    end

    # If `isinf(h)==true`, we use the maximum (finite)
    # step-size obtained from all coefficients as above.
    # Note that the time step is independent from `epsilon`.
    if isinf(h)
        h = zero(R)
        for i in eachindex(q)
            @inbounds hi = _second_stepsize(q[i], epsilon)
            h = max( h, hi )
        end
    end
    return h::R
end

"""
    _stepsize(aux1, epsilon, k)

Helper function to avoid code repetition.
Returns ``(epsilon/aux1)^(1/k)``.
"""
@inline function _stepsize(aux1::U, epsilon::T, k::Int) where {T<:Real, U<:Number}
    aux = epsilon / aux1
    kinv = 1 / k
    return aux^kinv
end

"""
    _second_stepsize(x, epsilon)

Corresponds to the "second stepsize control" in Jorba and Zou
(2005) paper. We use it if [`stepsize`](@ref) returns `Inf`.
"""
function _second_stepsize(x::Taylor1{U}, epsilon::T) where {T<:Real, U<:Number}
    R = promote_type(typeof(norm(constant_term(x), Inf)), T)
    x == zero(x) && return convert(R, Inf)
    ord = x.order
    z = zero(R)
    u = one(R)
    h = z
    for k in 1:ord-2
        @inbounds aux = norm( x[k], Inf)
        aux == z && continue
        aux1 = _stepsize(aux, u, k)
        h = max(h, aux1)
    end
    return h::R
end

#taylorstep
@doc doc"""
    taylorstep!(f, t, x, t0, order, abstol, params, tmpTaylor, arrTaylor, parse_eqs=true) -> δt
    taylorstep!(f!, t, x, dx, xaux, t0, order, abstol, params, tmpTaylor, arrTaylor, parse_eqs=true) -> δt

One-step Taylor integration for the one-dependent variable ODE ``\dot{x}=dx/dt=f(x, p, t)``
with initial conditions ``x(t_0)=x_0``.
Returns the time-step `δt` of the actual integration carried out (δt is positive).

Here, `f` is the function defining the RHS of the ODE (see
[`taylorinteg`](@ref)), `t` is the
independent variable, `x` contains the Taylor expansion of the dependent
variable, `order` is
the degree  used for the `Taylor1` polynomials during the integration
`abstol` is the absolute tolerance used to determine the time step
of the integration, and `params` are the parameters entering the ODE
functions.
For several variables, `dx` and `xaux`, both of the same type as `x`,
are needed to save allocations. Finally, `parse_eqs` is a switch
to force *not* using (`parse_eqs=false`) the specialized method of `jetcoeffs!`
created with [`@taylorize`](@ref); the default is `true` (parse the equations).
Finally, `parse_eqs` is a switch
to force *not* using (`parse_eqs=false`) the specialized method of `jetcoeffs!`
created with [`@taylorize`](@ref); the default is `true` (parse the equations).

"""
function taylorstep!(f, t::Taylor1{T}, x::Taylor1{U}, abstol::T, params,
        tmpTaylor, arrTaylor, parse_eqs::Bool=true) where {T<:Real, U<:Number}

    # Compute the Taylor coefficients
    __jetcoeffs!(Val(parse_eqs), f, t, x, params, tmpTaylor, arrTaylor)

    # Compute the step-size of the integration using `abstol`
    δt = stepsize(x, abstol)
    if isinf(δt)
        δt = _second_stepsize(x, abstol)
    end

    return δt
end

function taylorstep!(f!, t::Taylor1{T}, x::Vector{Taylor1{U}},
        dx::Vector{Taylor1{U}}, xaux::Vector{Taylor1{U}}, abstol::T, params,
        tmpTaylor, arrTaylor, parse_eqs::Bool=true) where {T<:Real, U<:Number}

    # Compute the Taylor coefficients
    __jetcoeffs!(Val(parse_eqs), f!, t, x, dx, xaux, params, tmpTaylor, arrTaylor)

    # Compute the step-size of the integration using `abstol`
    δt = stepsize(x, abstol)

    return δt
end

"""
    _determine_parsing!(parse_eqs::Bool, f, t, x, params)
    _determine_parsing!(parse_eqs::Bool, f, t, x, dx, params)

Check if the parsed method of `jetcoeffs!` exists and check it
runs without error.
"""
function _determine_parsing!(parse_eqs::Bool, f, t, x, params)
    parse_eqs = parse_eqs && !isempty(methodswith(Val{f}, jetcoeffs!)) &&
                             !isempty(methodswith(Val{f}, _allocate_jetcoeffs!))

    tmpTaylor, arrTaylor = _allocate_jetcoeffs!(t, x, params)

    if parse_eqs
        try
            tmpTaylor, arrTaylor = _allocate_jetcoeffs!(Val(f), t, x, params)
            jetcoeffs!(Val(f), t, x, params, tmpTaylor, arrTaylor)
        catch
            @warn("""Unable to use the parsed method of `jetcoeffs!` for `$f`,
            despite of having `parse_eqs=true`, due to some internal error.
            Using `parse_eqs = false`.""")
            parse_eqs = false
        end
    end

    return parse_eqs, tmpTaylor, arrTaylor
end
function _determine_parsing!(parse_eqs::Bool, f, t, x, dx, params)
    parse_eqs = parse_eqs && !isempty(methodswith(Val{f}, jetcoeffs!)) &&
                             !isempty(methodswith(Val{f}, _allocate_jetcoeffs!))

    tmpTaylor, arrTaylor = _allocate_jetcoeffs!(t, x, dx, params)

    if parse_eqs
        try
            tmpTaylor, arrTaylor = _allocate_jetcoeffs!(Val(f), t, x, dx, params)
            jetcoeffs!(Val(f), t, x, dx, params, tmpTaylor, arrTaylor)
        catch
            @warn("""Unable to use the parsed method of `jetcoeffs!` for `$f`,
            despite of having `parse_eqs=true`, due to some internal error.
            Using `parse_eqs = false`.""")
            parse_eqs = false
        end
    end

    return parse_eqs, tmpTaylor, arrTaylor
end

# taylorinteg
@doc doc"""
    taylorinteg(f, x0, t0, tmax, order, abstol, params[=nothing]; kwargs... )
    taylorinteg(f, x0, t0, tmax, order, abstol, Val(true), params[=nothing]; kwargs... )

General-purpose Taylor integrator for the explicit ODE ``\dot{x}=f(x, p, t)``.
The initial condition are specified by `x0` at time `t0`, and any parameters
encoded in `params`. The initial condition `x0` may be of type `T<:Number`
or `Vector{T}`, with `T` including `TaylorN{T}`; the latter case
is of interest for [jet transport applications](@ref jettransport).

The equations of motion are specified by the function `f`; we follow the same
convention of `DifferentialEquations.jl` to define this function, i.e.,
`f(x, p, t)` or `f!(dx, x, p, t)`; see the examples below.

It returns a vector with the values of time (independent variable),
and a vector with the computed values of
the dependent variable(s), and if the method used involves `Val(true)` it also
outputs the Taylor polynomial solutions obtained at each time step.
The integration stops when time is larger than `tmax`, in which case the last returned
value(s) correspond to that time, or when the number of saved steps is larger
than `maxsteps`.

The integration method uses polynomial expansions on the independent variable
of order `order`; the parameter `abstol` serves to define the
time step using the last two Taylor coefficients of the expansions.
Make sure you use a *large enough* `order` to assure convergence.

Currently, the recognized keyword arguments are:
- `maxsteps[=500]`: maximum number of integration steps.
- `parse_eqs[=true]`: use the specialized method of `jetcoeffs!` created
    with [`@taylorize`](@ref).

## Examples

For one dependent variable the function `f` defines the RHS of the equation of
motion, returning the value of ``\dot{x}``. The arguments of
this function are `(x, p, t)`, where `x` are the dependent variables, `p` are
the paremeters and `t` is the independent variable.

For several (two or more) dependent variables, the function `f!` defines
the RHS of the equations of motion, mutating (in-place) the (preallocated) vector
with components of ``\dot{x}``. The arguments of this function are `(dx, x, p, t)`,
where `dx` is the preallocated vector of ``\dot{x}``, `x` are the dependent
variables, `p` are the paremeters entering the ODEs and `t` is the independent
variable. The function may return this vector or simply `nothing`.

```julia
using TaylorIntegration

f(x, p, t) = x^2

tv, xv = taylorinteg(f, 3, 0.0, 0.3, 25, 1.0e-20, maxsteps=100 )

function f!(dx, x, p, t)
    for i in eachindex(x)
        dx[i] = x[i]^2
    end
    return nothing
end

tv, xv = taylorinteg(f!, [3, 3], 0.0, 0.3, 25, 1.0e-20, maxsteps=100 )

tv, xv, polynv = taylorinteg(f!, [3, 3], 0.0, 0.3, 25, 1.0e-20, maxsteps=100, Val(true) )
```

"""
function taylorinteg(f, x0::U, t0::T, tmax::T, order::Int, abstol::T,
    params = nothing; maxsteps::Int=500, parse_eqs::Bool=true) where {T<:Real, U<:Number}

    # Allocation
    tv = Array{T}(undef, maxsteps+1)
    xv = Array{U}(undef, maxsteps+1)

    # Initialize the Taylor1 expansions
    t = Taylor1( T, order )
    x = Taylor1( x0, order )

    # Initial conditions
    nsteps = 1
    @inbounds t[0] = t0
    @inbounds tv[1] = t0
    @inbounds xv[1] = x0
    sign_tstep = copysign(1, tmax-t0)

    # Determine if specialized jetcoeffs! method exists
    parse_eqs, tmpTaylor, arrTaylor = _determine_parsing!(parse_eqs, f, t, x, params)

    # Integration
    while sign_tstep*t0 < sign_tstep*tmax
        δt = taylorstep!(f, t, x, abstol, params, tmpTaylor, arrTaylor, parse_eqs) # δt is positive!
        # Below, δt has the proper sign according to the direction of the integration
        δt = sign_tstep * min(δt, sign_tstep*(tmax-t0))
        x0 = evaluate(x, δt) # new initial condition
        @inbounds x[0] = x0
        t0 += δt
        @inbounds t[0] = t0
        nsteps += 1
        @inbounds tv[nsteps] = t0
        @inbounds xv[nsteps] = x0
        if nsteps > maxsteps
            @warn("""
            Maximum number of integration steps reached; exiting.
            """)
            break
        end
    end

    #return tv, xv
    return view(tv,1:nsteps), view(xv,1:nsteps)
end

function taylorinteg(f, x0::U, t0::T, tmax::T, order::Int, abstol::T,
        ::Val{true}, params = nothing;
        maxsteps::Int=500, parse_eqs::Bool=true) where {T<:Real, U<:Number}

    # Allocation
    tv = Array{T}(undef, maxsteps+1)
    xv = Array{U}(undef, maxsteps+1)
    polynV = Array{Taylor1{U}}(undef, maxsteps+1)

    # Initialize the Taylor1 expansions
    t = Taylor1( T, order )
    x = Taylor1( x0, order )

    # Initial conditions
    nsteps = 1
    @inbounds t[0] = t0
    @inbounds tv[1] = t0
    @inbounds xv[1] = x0
    @inbounds polynV[1] = deepcopy(x)
    sign_tstep = copysign(1, tmax-t0)

    # Determine if specialized jetcoeffs! method exists
    parse_eqs, tmpTaylor, arrTaylor = _determine_parsing!(parse_eqs, f, t, x, params)

    # Integration
    while sign_tstep*t0 < sign_tstep*tmax
        δt = taylorstep!(f, t, x, abstol, params, tmpTaylor, arrTaylor, parse_eqs) # δt is positive!
        # Below, δt has the proper sign according to the direction of the integration
        δt = sign_tstep * min(δt, sign_tstep*(tmax-t0))
        x0 = evaluate(x, δt) # new initial condition
        @inbounds polynV[nsteps+1] = deepcopy(x)
        @inbounds x[0] = x0
        t0 += δt
        @inbounds t[0] = t0
        nsteps += 1
        @inbounds tv[nsteps] = t0
        @inbounds xv[nsteps] = x0
        if nsteps > maxsteps
            @warn("""
            Maximum number of integration steps reached; exiting.
            """)
            break
        end
    end

    return view(tv,1:nsteps), view(xv,1:nsteps), view(polynV, 1:nsteps)
end

function taylorinteg(f!, q0::Array{U,1}, t0::T, tmax::T, order::Int, abstol::T,
        params = nothing; maxsteps::Int=500, parse_eqs::Bool=true) where {T<:Real, U<:Number}

    # Allocation
    tv = Array{T}(undef, maxsteps+1)
    dof = length(q0)
    xv = Array{U}(undef, dof, maxsteps+1)

    # Initialize the vector of Taylor1 expansions
    t = Taylor1(T, order)
    x = Array{Taylor1{U}}(undef, dof)
    dx = Array{Taylor1{U}}(undef, dof)
    xaux = Array{Taylor1{U}}(undef, dof)
    for i in eachindex(q0)
        @inbounds x[i] = Taylor1( q0[i], order )
        @inbounds dx[i] = Taylor1( zero(q0[i]), order )
    end

    # Initial conditions
    @inbounds t[0] = t0
    x .= Taylor1.(q0, order)
    x0 = deepcopy(q0)
    @inbounds tv[1] = t0
    @inbounds xv[:,1] .= q0
    sign_tstep = copysign(1, tmax-t0)

    # Determine if specialized jetcoeffs! method exists
    parse_eqs, tmpTaylor, arrTaylor = _determine_parsing!(parse_eqs, f!, t, x, dx, params)

    # Integration
    nsteps = 1
    while sign_tstep*t0 < sign_tstep*tmax
        δt = taylorstep!(f!, t, x, dx, xaux, abstol, params, tmpTaylor, arrTaylor, parse_eqs) # δt is positive!
        # Below, δt has the proper sign according to the direction of the integration
        δt = sign_tstep * min(δt, sign_tstep*(tmax-t0))
        evaluate!(x, δt, x0) # new initial condition
        for i in eachindex(x0)
            @inbounds x[i][0] = x0[i]
            @inbounds dx[i] = Taylor1( zero(x0[i]), order )
        end
        t0 += δt
        @inbounds t[0] = t0
        nsteps += 1
        @inbounds tv[nsteps] = t0
        @inbounds xv[:,nsteps] .= x0
        if nsteps > maxsteps
            @warn("""
            Maximum number of integration steps reached; exiting.
            """)
            break
        end
    end

    return view(tv,1:nsteps), view(transpose(view(xv,:,1:nsteps)),1:nsteps,:)
end

# Method to also return (output) the Taylor polynomials of the integration
function taylorinteg(f!, q0::Array{U,1}, t0::T, tmax::T, order::Int, abstol::T,
        ::Val{true}, params = nothing;
        maxsteps::Int=500, parse_eqs::Bool=true) where {T<:Real, U<:Number}

    # Allocation
    tv = Array{T}(undef, maxsteps+1)
    dof = length(q0)
    xv = Array{U}(undef, dof, maxsteps+1)
    polynV = Array{Taylor1{U}}(undef, dof, maxsteps+1)

    # Initialize the vector of Taylor1 expansions
    t = Taylor1(T, order)
    x = Array{Taylor1{U}}(undef, dof)
    dx = Array{Taylor1{U}}(undef, dof)
    xaux = Array{Taylor1{U}}(undef, dof)
    for i in eachindex(q0)
        @inbounds x[i] = Taylor1( q0[i], order )
        @inbounds dx[i] = Taylor1( zero(q0[i]), order )
    end

    # Initial conditions
    @inbounds t[0] = t0
    x .= Taylor1.(q0, order)
    x0 = deepcopy(q0)
    @inbounds tv[1] = t0
    @inbounds xv[:,1] .= q0
    @inbounds polynV[:,1] .= deepcopy.(x)
    sign_tstep = copysign(1, tmax-t0)

    # Determine if specialized jetcoeffs! method exists
    parse_eqs, tmpTaylor, arrTaylor = _determine_parsing!(parse_eqs, f!, t, x, dx, params)

    # Integration
    nsteps = 1
    while sign_tstep*t0 < sign_tstep*tmax
        δt = taylorstep!(f!, t, x, dx, xaux, abstol, params, tmpTaylor, arrTaylor, parse_eqs) # δt is positive!
        # Below, δt has the proper sign according to the direction of the integration
        δt = sign_tstep * min(δt, sign_tstep*(tmax-t0))
        evaluate!(x, δt, x0) # new initial condition
        # Store the Taylor polynomial solution
        @inbounds polynV[:,nsteps+1] .= deepcopy.(x)
        for i in eachindex(x0)
            @inbounds x[i][0] = x0[i]
            @inbounds dx[i] = Taylor1( zero(x0[i]), order )
        end
        t0 += δt
        @inbounds t[0] = t0
        nsteps += 1
        @inbounds tv[nsteps] = t0
        @inbounds xv[:,nsteps] .= x0
        if nsteps > maxsteps
            @warn("""
            Maximum number of integration steps reached; exiting.
            """)
            break
        end
    end

    return view(tv,1:nsteps), view(transpose(view(xv,:,1:nsteps)),1:nsteps,:),
        view(transpose(view(polynV, :, 1:nsteps)), 1:nsteps, :)
end

# Integrate and return results evaluated at given time
@doc doc"""
    taylorinteg(f, x0, trange, order, abstol, params[=nothing]; keyword... )

General-purpose Taylor integrator for the explicit ODE
``\dot{x}=f(t,x)`` with initial condition specified by `x0::{T<:Number}`
or `x0::Vector{T}` at time `t0`.

The method returns a vector with the computed values of
the dependent variable(s), evaluated *only* at the times specified by
the range `trange`.

## Examples

```julia
xv = taylorinteg(f, 3, 0.0:0.001:0.3, 25, 1.0e-20, maxsteps=100 )

xv = taylorinteg(f!, [3, 3], 0.0:0.001:0.3, 25, 1.0e-20, maxsteps=100 );

```

"""
function taylorinteg(f, x0::U, trange::AbstractVector{T},
        order::Int, abstol::T, params = nothing;
        maxsteps::Int=500, parse_eqs::Bool=true) where {T<:Real, U<:Number}

    # Check if trange is increasingly or decreasingly sorted
    @assert (issorted(trange) ||
        issorted(reverse(trange))) "`trange` or `reverse(trange)` must be sorted"

    # Allocation
    nn = length(trange)
    xv = Array{U}(undef, nn)
    fill!(xv, T(NaN))

    # Initialize the Taylor1 expansions
    t = Taylor1( T, order )
    x = Taylor1( x0, order )

    # Initial conditions
    @inbounds t0, t1, tmax = trange[1], trange[2], trange[end]
    sign_tstep = copysign(1, tmax-t0)
    @inbounds t[0] = t0
    @inbounds xv[1] = x0

    # Determine if specialized jetcoeffs! method exists
    parse_eqs, tmpTaylor, arrTaylor = _determine_parsing!(parse_eqs, f, t, x, params)

    # Integration
    iter = 2
    nsteps = 1
    while sign_tstep*t0 < sign_tstep*tmax
        δt = taylorstep!(f, t, x, abstol, params, tmpTaylor, arrTaylor, parse_eqs)# δt is positive!
        # Below, δt has the proper sign according to the direction of the integration
        δt = sign_tstep * min(δt, sign_tstep*(tmax-t0))
        x0 = evaluate(x, δt) # new initial condition
        tnext = t0+δt
        # Evaluate solution at times within convergence radius
        while sign_tstep*t1 < sign_tstep*tnext
            x1 = evaluate(x, t1-t0)
            @inbounds xv[iter] = x1
            iter += 1
            @inbounds t1 = trange[iter]
        end
        if δt == tmax-t0
            @inbounds xv[iter] = x0
            break
        end
        @inbounds x[0] = x0
        t0 = tnext
        @inbounds t[0] = t0
        nsteps += 1
        if nsteps > maxsteps
            @warn("""
            Maximum number of integration steps reached; exiting.
            """)
            break
        end
    end
    return xv
end

function taylorinteg(f!, q0::Array{U,1}, trange::AbstractVector{T},
        order::Int, abstol::T, params = nothing; maxsteps::Int=500,
        parse_eqs::Bool=true) where {T<:Real, U<:Number}

    # Check if trange is increasingly or decreasingly sorted
    @assert (issorted(trange) ||
        issorted(reverse(trange))) "`trange` or `reverse(trange)` must be sorted"

    # Allocation
    nn = length(trange)
    dof = length(q0)
    x0 = similar(q0, eltype(q0), dof)
    x1 = similar(x0)
    fill!(x0, T(NaN))
    xv = Array{eltype(q0)}(undef, dof, nn)
    for ind in 1:nn
        @inbounds xv[:,ind] .= x0
    end

    # Initialize the vector of Taylor1 expansions
    t = Taylor1( T, order )
    x = Array{Taylor1{U}}(undef, dof)
    dx = Array{Taylor1{U}}(undef, dof)
    xaux = Array{Taylor1{U}}(undef, dof)
    for i in eachindex(q0)
        @inbounds x[i] = Taylor1( q0[i], order )
        @inbounds dx[i] = Taylor1( zero(q0[i]), order )
    end

    # Initial conditions
    @inbounds t[0] = trange[1]
    @inbounds t0, t1, tmax = trange[1], trange[2], trange[end]
    sign_tstep = copysign(1, tmax-t0)
    x .= Taylor1.(q0, order)
    @inbounds x0 .= q0
    @inbounds xv[:,1] .= q0

    # Determine if specialized jetcoeffs! method exists
    parse_eqs, tmpTaylor, arrTaylor = _determine_parsing!(parse_eqs, f!, t, x, dx, params)

    # Integration
    iter = 2
    nsteps = 1
    while sign_tstep*t0 < sign_tstep*tmax
        δt = taylorstep!(f!, t, x, dx, xaux, abstol, params, tmpTaylor, arrTaylor, parse_eqs) # δt is positive!
        # Below, δt has the proper sign according to the direction of the integration
        δt = sign_tstep * min(δt, sign_tstep*(tmax-t0))
        evaluate!(x, δt, x0) # new initial condition
        tnext = t0+δt
        # Evaluate solution at times within convergence radius
        while sign_tstep*t1 < sign_tstep*tnext
            evaluate!(x, t1-t0, x1)
            @inbounds xv[:,iter] .= x1
            iter += 1
            @inbounds t1 = trange[iter]
        end
        if δt == tmax-t0
            @inbounds xv[:,iter] .= x0
            break
        end
        for i in eachindex(x0)
            @inbounds x[i][0] = x0[i]
            @inbounds dx[i] = Taylor1( zero(x0[i]), order )
        end
        t0 = tnext
        @inbounds t[0] = t0
        nsteps += 1
        if nsteps > maxsteps
            @warn("""
            Maximum number of integration steps reached; exiting.
            """)
            break
        end
    end

    return transpose(xv)
end

# Generic functions
for R in (:Number, :Integer)
    @eval begin
        function taylorinteg(f, xx0::S, tt0::T, ttmax::U, order::Int, aabstol::V,
                params = nothing; maxsteps::Int=500, parse_eqs::Bool=true) where
                    {S<:$R, T<:Real, U<:Real, V<:Real}

            # In order to handle mixed input types, we promote types before integrating:
            t0, tmax, abstol, bfloat = promote(tt0, ttmax, aabstol, one(Float64))
            x0, tfloat1 = promote(xx0, t0)

            taylorinteg(f, x0, t0, tmax, order, abstol, params, maxsteps=maxsteps,
                parse_eqs=parse_eqs)
        end

        function taylorinteg(f, q0::Array{S,1}, tt0::T, ttmax::U, order::Int, aabstol::V,
                params = nothing; maxsteps::Int=500, parse_eqs::Bool=true) where
                    {S<:$R, T<:Real, U<:Real, V<:Real}

            #promote to common type before integrating:
            t0, tmax, abstol, afloat = promote(tt0, ttmax, aabstol, one(Float64))
            elq0, tt0 = promote(q0[1], t0)
            #convert the elements of q0 to the common, promoted type:
            q0_ = convert(Array{typeof(elq0)}, q0)

            taylorinteg(f, q0_, t0, tmax, order, abstol, params, maxsteps=maxsteps,
                parse_eqs=parse_eqs)
        end

        function taylorinteg(f, xx0::S, trange::AbstractVector{T},
                order::Int, aabstol::U, params = nothing;
                maxsteps::Int=500, parse_eqs::Bool=true) where
                    {S<:$R, T<:Real, U<:Real}
                    #
                t0, abstol, bfloat = promote(trange[1], aabstol, one(Float64))
                x0, tfloat1 = promote(xx0, t0)
                taylorinteg(f, x0, vec(trange.*one(t0)), order, abstol,
                    params, maxsteps=maxsteps, parse_eqs=parse_eqs)
        end

        function taylorinteg(f, q0::Array{S,1}, trange::AbstractVector{T},
                order::Int, aabstol::U, params = nothing;
                maxsteps::Int=500, parse_eqs::Bool=true) where
                    {S<:$R, T<:Real, U<:Real}
                    #
                t0, abstol, bfloat = promote(trange[1], aabstol, one(Float64))
                elq0, tt0 = promote(q0[1], t0)
                q0_ = convert(Array{typeof(elq0)}, q0)

                taylorinteg(f, q0_, vec(trange.*one(t0)), order, abstol,
                    params, maxsteps=maxsteps, parse_eqs=parse_eqs)
        end

        function taylorinteg(f, xx0::S, tt0::T, ttmax::U, order::Int, aabstol::V,
                ::Val{true}, params = nothing; maxsteps::Int=500, parse_eqs::Bool=true) where
                    {S<:$R, T<:Real, U<:Real, V<:Real}

            # In order to handle mixed input types, we promote types before integrating:
            t0, tmax, abstol, bfloat = promote(tt0, ttmax, aabstol, one(Float64))
            x0, tfloat1 = promote(xx0, t0)

            return taylorinteg(f, x0, t0, tmax, order, abstol, Val(true), params,
                maxsteps=maxsteps, parse_eqs=parse_eqs)
        end

        function taylorinteg(f, q0::Array{S,1}, tt0::T, ttmax::U, order::Int, aabstol::V,
                ::Val{true}, params = nothing; maxsteps::Int=500, parse_eqs::Bool=true) where
                    {S<:$R, T<:Real, U<:Real, V<:Real}

            #promote to common type before integrating:
            t0, tmax, abstol, afloat = promote(tt0, ttmax, aabstol, one(Float64))
            elq0, tt0 = promote(q0[1], t0)
            #convert the elements of q0 to the common, promoted type:
            q0_ = convert(Array{typeof(elq0)}, q0)

            return taylorinteg(f, q0_, t0, tmax, order, abstol, Val(true), params,
                maxsteps=maxsteps, parse_eqs=parse_eqs)
        end
    end
end


function itr2bin(x)
    res = 0
    for v in x
        res += v
        res <<= 1
    end
    return res >> 1
end

function itr2ind(x, sizes)
    n = length(sizes)
    res = 0
    _mul = 1
    for i in n : -1 : 1
        res += x[i] * _mul
        _mul *= sizes[i]
    end
    return res
end

# fill array v with bits of number x
function fill_bits!(v, x)
    for i in length(v) : -1 : 1
        v[i] = x & 1
        x >>= 1
    end
end

# fill array v with digits of number x in base-system
function fill_numbers!(v, x; base :: Int)
    for i in length(v) : -1 : 1
        v[i] = x % base
        x = div(x, base)
    end
end

# set bits with numbers vs[1],
# ..., vs[k] to 1
# return 2^(vs[1] - 1) + ... + 2^(vs[k] - 1)
function set_bits(vs)
    res = 0
    for num in vs
        res += 2^(num - 1)
    end
    return res
end


# # Network Matrices

# `PowerSystems.jl` is able to build classic power systems modeling network matrices such as
# [Ybus](https://en.wikipedia.org/wiki/Nodal_admittance_matrix), [PTDF](https://www.powerworld.com/WebHelp/Content/MainDocumentation_HTML/Power_Transfer_Distribution_Factors.htm) and LODF.

# Check section [Network Matrices](@ref net_mat) for more details.

# ## Overview
#
# Network matrices are implemented in `PowerSystems.jl` as arrays that enable using Branch or Bus
# names as indexes to facilitate exploration and analysis. Ybus is stored as an SparseMatrix
# and the PTDF and LODF are stored as dense matrices. **Note*** Ybus is converted to a dense
# matrix for printing in the REPL,
# The network matrices code implements the Goderya algorithm to find islands.

using PowerSystems
const PSY = PowerSystems
DATA_DIR = "../../../data" #hide
system_data = System(joinpath(DATA_DIR, "matpower/case14.m"))

# ## Ybus
#
# The ybus can be calculated as follows:
ybus = Ybus(system_data)

# The matrix can be indexed using directly the bus numbers. In this example buses are numbered
# 1-14. However, in large systems buses don't usually follow sequential numbering. You can access
# the entries of the Ybus with direct indexing or using the buses.

ybus_entry = ybus[3,3]
#
bus3 = get_component(Bus, system_data, "Bus 3     HV")
ybus_entry = ybus[bus3, bus3]

# We recognize that many models require matrix operations. For those cases, you can access the
# sparse data as follows:

sparse_array = get_data(ybus)

# ## PTDF
#
# The PTDF matrix can be calculated as follows:
ptdf = PTDF(system_data)

# The entries of the matrix can also be indexed using the devices directly:
line = get_component(Line, system_data, "Bus 1     HV-Bus 2     HV-i_1")
bus = get_component(Bus, system_data, "Bus 3     HV")

ptdf_entry = ptdf[line, bus]

# Alternatively, the branch name and bus number can be used. For instance, for the same PTDF
# entry we can do:

ptdf_entry = ptdf["Bus 1     HV-Bus 2     HV-i_1", 3]

# PTDF also takes a vector of distributed slacks, for now this feature requires passing a
# vector of weights with the same number of elements are buses in the system. For more details
# check the API entry for [`PTDF`](@ref).

# In the same fashion as in the case of the Ybus, we recognize many models use matrix operations
# when using the PTDF matrix. The raw data is available using `get_data` and the lookup table
# between (names,numbers) and (i,j) matrix entries is available using [`get_lookup`](@ref)

branch_lookup, bus_lookup = get_lookup(ptdf);

# Branch lookup:

branch_lookup

# Bus Look up

bus_lookup

# ## LODF

# LODF and PTDF share the same characteristics in terms of indexing and calculation:

lodf = LODF(system_data)


# filled contours (contourf)

using CairoMakie

xs = LinRange(0, 10, 100)
ys = LinRange(0, 10, 100)
zs = [cos(x) * sin(y) for x in xs, y in ys]

f = Figure()
Axis(f[1, 1])

co = CairoMakie.contourf!(xs, ys, zs, levels = 10, colormap = :flag_ar)

Colorbar(f[1, 2], co)