text:
# MIT License
#
# Copyright (c) 2018 Martin Biel
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

@everywhere module TestDecisionExpressions

using StochasticPrograms
using Test
using MathOptInterface
const MOI = MathOptInterface
const MOIU = MOI.Utilities
const MA = MOI.MutableArithmetics

function DecisionModel()
    m = Model()
    m.ext[:decisions] = Decisions(Val{1}())
    return m
end

macro decision_variable(m, x)
    return esc(quote
        @variable($m, $x, set = StochasticPrograms.DecisionSet(1))
    end)
end

macro test_expression(expr)
    return esc(
        quote
            @test JuMP.isequal_canonical(@expression(model, $expr), $expr)
        end,
    )
end

macro test_expression_with_string(expr, str)
    return esc(
        quote
            realized_expr = @inferred $expr
            @test string(realized_expr) == $str
            @test JuMP.isequal_canonical(
                @expression(model, $expr),
                realized_expr,
            )
        end,
    )
end

function test_affine_expressions()
    @testset "isequal(::DecisionAffExpr)" begin
        m = DecisionModel()
        @decision_variable(m, x)
        @variable(m, y)
        @test isequal(x + y + 1, x + y + 1)
    end

    @testset "hash(::DecisionAffExpr)" begin
        m = DecisionModel()
        @decision_variable(m, x)
        @variable(m, y)
        @test hash(x + y + 1) == hash(x + y + 1)
    end

    @testset "drop_zeros!(::DecisionAffExpr)" begin
        m = DecisionModel()
        @decision_variable(m, x[1:2])
        @variable(m, y[1:2])
        expr = x[1] + x[2] + y[1] + y[2] - x[2] - y[2] + 1
        @test !isequal(expr, x[1] + y[1] + 1)
        JuMP.drop_zeros!(expr)
        @test isequal(expr, x[1] + y[1] + 1)
    end

    @testset "iszero(::DecisionAffExpr)" begin
        m = DecisionModel()
        @decision_variable(m, x)
        @variable(m, y)
        @test !iszero(x + y + 1)
        @test !iszero(x + y + 0)
        @test iszero(0 * x + 0 * y + 0)
        @test iszero(x - x + y - y)
    end

    @testset "add_to_expression!(::DecisionAffExpr{C}, ::VariableRef/DecisionRef)" begin
        m = DecisionModel()
        @decision_variable(m, x)
        @variable(m, y)
        aff = DecisionAffExpr(
            JuMP.GenericAffExpr(1.0, y => 2.0),
            zero(JuMP.GenericAffExpr{Float64,DecisionRef}),
        )
        @test JuMP.isequal_canonical(
            JuMP.add_to_expression!(aff, x),
            DecisionAffExpr(
                JuMP.GenericAffExpr(1.0, y => 2.0),
                JuMP.GenericAffExpr(0.0, x => 1.0),
            ),
        )
        @test JuMP.isequal_canonical(
            JuMP.add_to_expression!(aff, y),
            DecisionAffExpr(
                JuMP.GenericAffExpr(1.0, y => 3.0),
                JuMP.GenericAffExpr(0.0, x => 1.0),
            ),
        )
    end

    @testset "add_to_expression!(::DecisionAffExpr{C}, ::C)" begin
        m = DecisionModel()
        @decision_variable(m, x)
        @variable(m, y)
        aff = DecisionAffExpr(
            JuMP.GenericAffExpr(1.0, y => 2.0),
            JuMP.GenericAffExpr(0.0, x => 2.0),
        )
        @test JuMP.isequal_canonical(
            JuMP.add_to_expression!(aff, 1.0),
            DecisionAffExpr(
                JuMP.GenericAffExpr(2.0, y => 2.0),
                JuMP.GenericAffExpr(0.0, x => 2.0),
            ),
        )
    end

    @testset "coefficient(aff::DecisionAffExpr, v::VariableRef/DecisionRef)" begin
        m = DecisionModel()
        x1 = @decision_variable(m, x1)
        x2 = @decision_variable(m, x2)
        y1 = @variable(m, y1)
        y2 = @variable(m, y2)
        aff = @expression(m, 1.0 * x1 + 1.0 * y1)
        @test coefficient(aff, x1) == 1.0
        @test coefficient(aff, y1) == 1.0
        @test coefficient(aff, x2) == 0.0
        @test coefficient(aff, y2) == 0.0
    end

    @testset "coefficient(aff::DecisionAffExpr, v1::VariableRef/DecisionRef, v2::VariableRef/DecisionRef)" begin
        m = DecisionModel()
        x = @decision_variable(m, x1)
        y = @variable(m, y1)
        aff = @expression(m, 1.0 * x + 1.0 * y)
        @test coefficient(aff, x, x) == 0.0
        @test coefficient(aff, x, y) == 0.0
        @test coefficient(aff, y, x) == 0.0
        @test coefficient(aff, y, y) == 0.0
    end

    @testset "(+)(::DecisionAffExpr)" begin
        model = DecisionModel()
        x = @decision_variable(model, x)
        y = @variable(model, y)
        @test_expression_with_string (+)(x + y + 1) "x + y + 1"
    end

    @testset "sum(::Vector{VariableRef/DecisionRef})" begin
        model = DecisionModel()
        x = @decision_variable(model, x[1:2])
        y = @variable(model, y[1:2])
        @test_expression_with_string sum(x) + sum(y) "x[1] + x[2] + y[1] + y[2]"
    end
end

function test_mutable_arithmetics()
    @testset "MA.add_mul!(ex::Number, c::Number, x::DecisionAffExpr)" begin
        m = DecisionModel()
        x = @decision_variable(m, x)
        y = @variable(m, y)
        aff = MA.add_mul!(1.0, 2.0, DecisionAffExpr(
            JuMP.GenericAffExpr(1.0, y => 1.0),
            JuMP.GenericAffExpr(0.0, x => 1.0),
        ))
        @test JuMP.isequal_canonical(aff, DecisionAffExpr(
            JuMP.GenericAffExpr(3.0, y => 2.0),
            JuMP.GenericAffExpr(0.0, x => 2.0),
        ))
    end

    @testset "MA.add_mul!(ex::Number, c::Number, x::DecisionQuadExpr) with c == 0" begin
        quad = MA.add_mul!(2.0, 0.0, DecisionQuadExpr{Float64}())
        @test JuMP.isequal_canonical(quad, convert(DecisionQuadExpr{Float64}, 2.0))
    end

    @testset "MA.add_mul!(ex::Number, c::VariableRef/DecisionRef, x::VariableRef/DecisionRef)" begin
        model = DecisionModel()
        x = @decision_variable(model, x)
        y = @variable(model, y)
        @test_expression_with_string MA.add_mul(5.0, x, x) "x² + 5"
        @test_expression_with_string MA.add_mul!(5.0, x, x) "x² + 5"
        @test_expression_with_string MA.add_mul(5.0, x, y) "x*y + 5"
        @test_expression_with_string MA.add_mul!(5.0, x, y) "x*y + 5"
        @test_expression_with_string MA.add_mul(5.0, y, x) "x*y + 5"
        @test_expression_with_string MA.add_mul!(5.0, y, x) "x*y + 5"

    end

    @testset "MA.add_mul!(ex::Number, c::T, x::T) where T<:DecisionAffExpr" begin
        model = DecisionModel()
        x = @decision_variable(model, x)
        y = @variable(model, y)
        @test_expression_with_string MA.add_mul(1.0, 2x + 2y, x + y + 1) "2 x² + 2 x + 4 x*y + 2 y² + 2 y + 1"
        @test_expression_with_string MA.add_mul!(1.0, 2x + 2y, x + y + 1) "2 x² + 2 x + 4 x*y + 2 y² + 2 y + 1"
    end

    @testset "MA.add_mul!(ex::Number, c::DecisionAffExpr{C}, x::VariableRef/DecisionRef) where C" begin
        model = DecisionModel()
        x = @decision_variable(model, x)
        y = @variable(model, y)
        @test_expression_with_string MA.add_mul(1.0, 2x + 2y, x) "2 x² + 2 x*y + 1"
        @test_expression_with_string MA.add_mul!(1.0, 2x + 2y, x) "2 x² + 2 x*y + 1"
        @test_expression_with_string MA.add_mul(1.0, 2x + 2y, y) "2 x*y + 2 y² + 1"
        @test_expression_with_string MA.add_mul!(1.0, 2x + 2y, y) "2 x*y + 2 y² + 1"
    end

    @testset "MA.add_mul!(ex::Number, c::DecisionQuadExpr, x::Number)" begin
        model = DecisionModel()
        x = @decision_variable(model, x)
        y = @variable(model, y)
        @test_expression_with_string MA.add_mul(0.0, x^2 + y^2 + x * y, 1.0) "x² + x*y + y²"
        @test_expression_with_string MA.add_mul!(0.0, x^2 + y^2 + x * y, 1.0) "x² + x*y + y²"
    end

    @testset "MA.add_mul!(ex::Number, c::DecisionQuadExpr, x::Number) with c == 0" begin
        model = DecisionModel()
        x = @decision_variable(model, x)
        y = @variable(model, y)
        @test_expression_with_string MA.add_mul(0.0, x^2 + y^2 + x * y, 0.0) "0"
        @test_expression_with_string MA.add_mul!(0.0, x^2 + y^2 + x * y, 0.0) "0"
    end

    @testset "MA.add_mul!(aff::DecisionAffExpr{C}, c::VariableRef, x::GenericAffExpr{C,VariableRef}) where C" begin
        model = DecisionModel()
        x = @decision_variable(model, x)
        y = @variable(model, y)
        @test_expression_with_string MA.add_mul(2x, y, y + 1) "2 x + y² + y"
        @test_expression_with_string MA.add_mul!(2x, y, y + 1) "2 x + y² + y"
    end

    @testset "MA.add_mul!(aff::DecisionAffExpr{C}, c::DecisionRef, x::GenericAffExpr{C,VariableRef}) where C" begin
        model = DecisionModel()
        x = @decision_variable(model, x)
        y = @variable(model, y)
        @test_expression_with_string MA.add_mul(2x, x, y + 1) "3 x + x*y"
        @test_expression_with_string MA.add_mul!(2x, x, y + 1) "3 x + x*y"
    end

    @testset "MA.add_mul!(aff::DecisionAffExpr, c::VariableRef, x::DecisionAffExpr)" begin
        model = DecisionModel()
        x = @decision_variable(model, x)
        y = @variable(model, y)
        @test_expression_with_string MA.add_mul(2x, y, x + 1) "2 x + x*y + y"
        @test_expression_with_string MA.add_mul!(2x, y, x + 1) "2 x + x*y + y"
    end

    @testset "MA.add_mul!(aff::DecisionAffExpr, c::DecisionRef, x::DecisionAffExpr)" begin
        model = DecisionModel()
        x = @decision_variable(model, x)
        y = @variable(model, y)
        @test_expression_with_string MA.add_mul(2x, x, x + y + 1) "x² + 3 x + x*y"
        @test_expression_with_string MA.add_mul!(2x, x, x + y + 1) "x² + 3 x + x*y"
    end

    @testset "MA.add_mul!(aff::DecisionAffExpr{C}, c::GenericAffExpr{C,VariableRef}, x::Number) where C" begin
        model = DecisionModel()
        x = @decision_variable(model, x)
        y = @variable(model, y)
        @test_expression_with_string MA.add_mul(2x, 2y, 1) "2 x + 2 y"
        @test_expression_with_string MA.add_mul!(2x, 2y, 1) "2 x + 2 y"
    end

    @testset "MA.add_mul!(aff::DecisionAffExpr{C}, c::GenericAffExpr{C,DecisionRef}, x::Number) where C" begin
        model = DecisionModel()
        x = @decision_variable(model, x)
        @test_expression_with_string MA.add_mul(2x, 2x, 1) "4 x"
        @test_expression_with_string MA.add_mul!(2x, 2x, 1) "4 x"
    end

    @testset "MA.add_mul!(aff::DecisionAffExpr{C}, c::GenericQuadExpr{C,VariableRef}, x::Number) where C" begin
        model = DecisionModel()
        x = @decision_variable(model, x)
        y = @variable(model, y)
        @test_expression_with_string MA.add_mul(2x, y^2, 1) "2 x + y²"
        @test_expression_with_string MA.add_mul!(2x, y^2, 1) "2 x + y²"
    end

    @testset "MA.add_mul!(aff::DecisionAffExpr{C}, c::DecisionQuadExpr{C}, x::Number) where C" begin
        model = DecisionModel()
        x = @decision_variable(model, x)
        y = @variable(model, y)
        @test_expression_with_string MA.add_mul(2x + 2y, x^2 + y^2, 1) "x² + 2 x + y² + 2 y"
        @test_expression_with_string MA.add_mul!(2x + 2y, x^2 + y^2, 1) "x² + 2 x + y² + 2 y"
    end

    @testset "MA.add_mul!(aff::DecisionAffExpr, c::DecisionQuadExpr, x::Number) with x == 0" begin
        model = DecisionModel()
        x = @decision_variable(model, x)
        y = @variable(model, y)
        @test_expression_with_string MA.add_mul(2x + 2y, x^2 + y^2, 0) "2 x + 2 y"
        @test_expression_with_string MA.add_mul!(2x + 2y, x^2 + y^2, 0) "2 x + 2 y"
    end

    @testset "MA.add_mul!(aff::DecisionQuadExpr, c::Number, x::DecisionAffExpr) with c == 0" begin
        model = DecisionModel()
        x = @decision_variable(model, x)
        y = @variable(model, y)
        @test_expression_with_string MA.add_mul(2x + 2y, 0, x^2 + y^2) "2 x + 2 y"
        @test_expression_with_string MA.add_mul!(2x + 2y, 0, x^2 + y^2) "2 x + 2 y"
    end

    @testset "MA.add_mul!(ex::DecisionAffExpr, c::DecisionAffExpr, x::DecisionAffExpr)" begin
        model = DecisionModel()
        x = @decision_variable(model, x)
        y = @variable(model, y)
        # GenericAffExpr, DecisionAffExpr
        @test_expression_with_string MA.add_mul(2x + 2y, y + 1, x + 0) "3 x + x*y + 2 y"
        @test_expression_with_string MA.add_mul!(2x + 2y, y + 1, x + 0) "3 x + x*y + 2 y"
        # DecisionAffExpr, GenericAffExpr
        @test_expression_with_string MA.add_mul(2x + 2y, x + 1, y + 0) "2 x + x*y + 3 y"
        @test_expression_with_string MA.add_mul!(2x + 2y, x + 1, y + 0) "2 x + x*y + 3 y"
        # DecisionAffExpr, DecisionAffExpr
        @test_expression_with_string MA.add_mul(2x + 2y, x + 1, x + 0) "x² + 3 x + 2 y"
        @test_expression_with_string MA.add_mul!(2x + 2y, x + 1, x + 0) "x² + 3 x + 2 y"
    end

    @testset "MA.add_mul!(quad::DecisionQuadExpr, c::DecisionAffExpr, x::Number)" begin
        model = DecisionModel()
        x = @decision_variable(model, x)
        y = @variable(model, y)
        # GenericAffExpr
        @test_expression_with_string MA.add_mul(x^2 + y^2, y + 1, 1) "x² + y² + y + 1"
        @test_expression_with_string MA.add_mul!(x^2 + y^2, y + 1, 1) "x² + y² + y + 1"
        # DecisionAffExpr
        @test_expression_with_string MA.add_mul(x^2 + y^2, x + y + 1, 1) "x² + x + y² + y + 1"
        @test_expression_with_string MA.add_mul!(x^2 + y^2, x + y + 1, 1) "x² + x + y² + y + 1"
    end

    @testset "MA.add_mul!(quad::DecisionQuadExpr, c::VariableRef, x::DecisionAffExpr)" begin
        model = DecisionModel()
        x = @decision_variable(model, x)
        y = @variable(model, y)
        @test_expression_with_string MA.add_mul(x^2 + y^2, y, x + 1) "x² + x*y + y² + y"
        @test_expression_with_string MA.add_mul!(x^2 + y^2, y, x + 1) "x² + x*y + y² + y"
    end

    @testset "MA.add_mul!(quad::DecisionQuadExpr, c::DecisionRef, x::DecisionAffExpr)" begin
        model = DecisionModel()
        x = @decision_variable(model, x)
        y = @variable(model, y)
        @test_expression_with_string MA.add_mul(x^2 + y^2, x, x + y + 1) "2 x² + x + x*y + y²"
        @test_expression_with_string MA.add_mul!(x^2 + y^2, x, x + y + 1) "2 x² + x + x*y + y²"
    end

    @testset "MA.add_mul!(quad::DecisionQuadExpr, c::DecisionQuadExpr, x::Number)" begin
        model = DecisionModel()
        x = @decision_variable(model, x)
        y = @variable(model, y)
        # GenericQuadExpr
        @test_expression_with_string MA.add_mul(x^2 + y^2 + x, y^2 + y, 2.0) "x² + x + 3 y² + 2 y"
        @test_expression_with_string MA.add_mul!(x^2 + y^2 + x, y^2 + y, 2.0) "x² + x + 3 y² + 2 y"
        # DecisionQuadExpr
        @test_expression_with_string MA.add_mul(x^2 + y^2 + x, x^2 + x, 2.0) "3 x² + 3 x + y²"
        @test_expression_with_string MA.add_mul!(x^2 + y^2 + x, x^2 + x, 2.0) "3 x² + 3 x + y²"
    end

    @testset "MA.add_mul!(ex::DecisionQuadExpr, c::DecisionAffExpr, x::DecisionAffExpr)" begin
        model = DecisionModel()
        x = @decision_variable(model, x)
        y = @variable(model, y)
        # GenericAffExpr, GenericAffExpr
        @test_expression_with_string MA.add_mul(x^2 + y^2 + x, y + 0, y + 1) "x² + x + 2 y² + y"
        @test_expression_with_string MA.add_mul!(x^2 + y^2 + x, y + 0, y + 1) "x² + x + 2 y² + y"
        # DecisionAffExpr, GenericAffExpr
        @test_expression_with_string MA.add_mul(x^2 + y^2 + x, x + 0, y + 1) "x² + 2 x + x*y + y²"
        @test_expression_with_string MA.add_mul!(x^2 + y^2 + x, x + 0, y + 1) "x² + 2 x + x*y + y²"
        # GenericAffExpr, DecisionAffExpr
        @test_expression_with_string MA.add_mul(x^2 + y^2 + x, y + 0, x + 1) "x² + x + x*y + y² + y"
        @test_expression_with_string MA.add_mul!(x^2 + y^2 + x, y + 0, x + 1) "x² + x + x*y + y² + y"
        # DecisionAffExpr, DecisionAffExpr
        @test_expression_with_string MA.add_mul(x^2 + y^2 + x, x + y, x + y + 1) "2 x² + 2 x + 2 x*y + 2 y² + y"
        @test_expression_with_string MA.add_mul!(x^2 + y^2 + x, x + y, x + y + 1) "2 x² + 2 x + 2 x*y + 2 y² + y"
    end
end

function test_quadratic_expressions()
    @testset "isequal(::DecisionQuadExpr)" begin
        m = DecisionModel()
        @decision_variable(m, x)
        @variable(m, y)
        @test isequal(x^2 + 2*x*y + y^2 + 1, x^2 + 2*x*y + y^2 + 1)
    end

    @testset "hash(::DecisionQuadExpr)" begin
        m = DecisionModel()
        @decision_variable(m, x)
        @variable(m, y)
        @test hash(x^2 + 2*x*y + y^2 + 1) == hash(x^2 + 2*x*y + y^2 + 1)
    end

    @testset "drop_zeros!(::DecisionQuadExpr)" begin
        m = DecisionModel()
        @decision_variable(m, x[1:2])
        @variable(m, y[1:2])
        expr = x[1]^2 + x[2]^2 + x[1] + x[2] + 2*x[1]*y[1] + 2*x[2]*y[2] + y[1]^2 + y[2]^2 + y[1] + y[2] - x[2]^2 - y[2]^2 - 2*x[2]*y[2] - x[2] - y[2] + 1
        @test !isequal(expr, x[1]^2 + x[1] + 2*x[1]*y[1] + y[1]^2 + y[1] + 1)
        JuMP.drop_zeros!(expr)
        @test isequal(expr, x[1]^2 + x[1] + 2*x[1]*y[1] + y[1]^2 + y[1] + 1)
    end

    @testset "iszero(::DecisionQuadExpr)" begin
        m = DecisionModel()
        @decision_variable(m, x)
        @variable(m, y)
        @test !iszero(x^2 + x*y + y^2 + 1)
        @test !iszero(x^2 + x*y + y^2 + 0)
        @test iszero(0 * x^2 + 0 * x + 0 * x * y + 0 * y^2 + 0 * y + 0)
        @test iszero(x^2 - x^2 + y^2 - y^2 + x*y - x*y)
    end

    @testset "coefficient(quad::DecisionQuadExpr, v::VariableRef/DecisionRef)" begin
        m = DecisionModel()
        x = @decision_variable(m, x)
        y = @variable(m, y)
        z = @decision_variable(m, z)
        w = @variable(m, w)
        quad = @expression(m, 6.0 * x^2 + 5.0 * x * y + 2.0 * y + 3.0 * x)
        @test coefficient(quad, x) == 3.0
        @test coefficient(quad, y) == 2.0
        @test coefficient(quad, z) == 0.0
        @test coefficient(quad, w) == 0.0
    end

    @testset "coefficient(quad::Quad, v1::VariableRef/DecisionRef, v2::VariableRef/DecisionRef)" begin
        m = DecisionModel()
        x = @decision_variable(m, x)
        y = @variable(m, y)
        z = @decision_variable(m, z)
        w = @variable(m, w)
        quad = @expression(m, 6.0 * x^2 + 5.0 * x * y + 4.0 * y^2 + 2.0 * y + 3.0 * x + y * w)
        @test coefficient(quad, x, y) == 5.0
        @test coefficient(quad, x, x) == 6.0
        @test coefficient(quad, y, y) == 4.0
        @test coefficient(quad, x, y) == coefficient(quad, y, x)
        @test coefficient(quad, y, w) == 1.0
        @test coefficient(quad, z, z) == 0.0
        @test coefficient(quad, w, w) == 0.0
    end

    @testset "(+)(::DecisionQuadExpr)" begin
        model = DecisionModel()
        x = @decision_variable(model, x)
        y = @variable(model, y)
        @test_expression_with_string (+)(x^2 + 2*x*y + y^2 + 1) "x² + 2 x*y + y² + 1"
    end
end

function runtests()
    @testset "DecisionExpressions" begin
        for name in names(@__MODULE__; all = true)
            if !startswith("$(name)", "test_")
                continue
            end
            f = getfield(@__MODULE__, name)
            @testset "$(name)" begin
                f()
            end
        end
    end
end

end


text:
export peturb

""""
    ensure_diagonal(Π::M)

If `M<:SparseMatrixCSC`, ensure that diagonal has all entries (even if zero).
"""
function ensure_diagonal(Π::M) where {T<:Number, M<:SparseMatrixCSC{T}}
    n = size(Π, 1)
    A = Π +  SparseMatrixCSC{T}(I, n, n)
    typeof(A) == M || throw(error("type of Πnew != typeof(Π)"))
    Arows = rowvals(A)
    Avals = nonzeros(A)
    for j = 1:n
       for i in nzrange(A, j)
          j == Arows[i]  &&  (Avals[i] -= one(T))
       end
    end
    return A
end

ensure_diagonal(Π::M) where {T<:Number, M<:Matrix{T}} = Π

peturb(x::T) where {T} = max( abs(x), one(T) ) * cbrt(eps(T))

ext: jl


text:
using LinearAlgebra
using SparseArrays
using Plots
using ForwardDiff

"This routine solves Poisson's equation using piecewise linear finite element methods.
The matrix is assembled locally, and boundary nodes are not removed."

m = 100 # number of elements

α = 1.0
β = 0.0
f(x) = x + 0*2*sin(2*pi*x)
κ(x) = 1.0 + x

# define spatial grid
x = LinRange(-1,1,m+2) # x_0, x_1, ..., x_m, x_{m+1} = x_0
xint = x[1:end-1]
h = x[2]-x[1]

# x = x + randn(size(x))/(2*m)
# x = @. x + (1+x)*(1-x)/3

# construct local FEM matrix
map_point(x,a,b) = a + (b-a) * (1+x)/2 # maps x ∈ [-1,1] to interval [a,b]
A = spzeros(m+2,m+2)
b = zeros(m+2)
A_local = [1 -1;-1 1]

function ϕ(r)
    ϕ1 = @. (1-r)/2
    ϕ2 = @. (1+r)/2
    return [ϕ1 ϕ2]
end
function dϕ(r)
    dϕ1 = @. -.5 + 0*r
    dϕ2 = @. .5 + 0*r
    return [dϕ1 dϕ2]
end

r = 0.0
w = 2.0

# r = [-1; 1] / sqrt(3)
# w = ones(2)
for e = 1:m+1
    h_e = x[e+1] - x[e]
    ids = e:e+1

    # use 1-point quadrature to approximate integrals
    x_e = map_point(0, x[e], x[e+1])

    # accumulate local contributions
    A[ids,ids] += A_local * κ(x_e) / h_e
    b[ids] += h_e * w * f(x_e) * [.5;.5]

    # # use 2-point Gauss quadrature to approximate integrals
    # x_e = map_point.(r, x[e], x[e+1])
    #
    # # build and accumulate local contributions
    # A_local = dϕ(r)'*diagm(w_e.*κ.(x_e))*dϕ(r)
    # A[ids,ids] += A_local / h_e
    # b[ids] += h_e * (ϕ(r)' * (w.*f.(x_e)))
end

# modify for boundary conditions
function impose_Dirichlet_BC!(A,b,i,val)
    b .-= A[:,i]*val
    b[i] = val
    A[i,:] .= 0
    A[:,i] .= 0
    A[i,i] = 1.0
end
function impose_Neumann_BC!(A,b,i,val)
    b[i] += val
end

impose_Dirichlet_BC!(A,b,1,α)
impose_Dirichlet_BC!(A,b,m+2,β)
# impose_Neumann_BC!(A,b,m+2,β)

u = A\b
plot(x,u,legend=false,mark=:dot,ms=2)


text:
#=
constraints:
- species groups e.g. as polytomies (type 1), or
- clades in the major tree (type 2), or
- clades in *some* displayed tree (type 3, not implemented).
outgroups "grades" can be coded as ingroup clades

move types, all under topological constraints:
- NNI = nearest neighbor interchange for undirected network,
  similar to NNIs for rooted networks defined in
  Gambette, van Iersel, Jones, Lafond, Pardi and Scornavacca (2017),
  Rearrangement moves on rooted phylogenetic networks.
  PLOS Computational Biology 13(8):e1005611.
- move the root: needed for optimization of a semi-directed network,
  because the root position affects the feasibility of the
  NNIs starting from a BR configuration (bifurcation -> reticulation).

- remove a hybrid edge: `deletehybridedge!` in file deleteHybrid.jl,
  but does not check for clade constraints
- add a hybrid edge: in file addHybrid.jl
- change the direction of a hybrid edge

in `moves_snaq.jl`, functions are tailored to level-1 networks;
here, functions apply to semidirected networks of all levels.
=#

"""
Type for various topological constraints, such as:

1. a set of taxa forming a clade in the major tree
2. a set of individuals belonging to the same species
3. a set of taxa forming a clade in any one of the displayed trees
"""
mutable struct TopologyConstraint
    "type of constraint: 1=species, 2=clade in major tree. 3=clade in some displayed tree: not yet implemented."
    type::UInt8
    "names of taxa in the constraint (clade / species members)"
    taxonnames::Vector{String}
    """
    node numbers of taxa in the constraint (clade / species members).
    warning: interpretation dependent on the internal network representation
    """
    taxonnums::Set{Int}
    "stem edge of the constrained group"
    edge::Edge
    "crown node, that is, child node of stem edge"
    node::Node
end

"""
    TopologyConstraint(type::UInt8, taxonnames::Vector{String}, net::HybridNetwork)

Create a topology constraint from user-given type, taxon names, and network.
There are 3 types of constraints:

- type 1: A set of tips that forms a species.
- type 2: A set of tips that forms a clade in the major tree.
  Note that the root matters. Constraining a set of species to be an outgroup
  is equivalent to constraining the ingroup to form a clade.
- type 3: A set of tips that forms a clade in any one of the displayed trees.

Note: currently, with type-1 constraints, hybridizations are prevented
from coming into or going out of a species group.

# examples

```jldoctest
julia> net = readTopology("(((2a,2b),(((((1a,1b,1c),4),(5)#H1),(#H1,(6,7))))#H2),(#H2,10));");

julia> c_species1 = PhyloNetworks.TopologyConstraint(0x01, ["1a","1b","1c"], net)
Species constraint, on tips: 1a, 1b, 1c
 stem edge number 7
 crown node number -9

julia> c_species2 = PhyloNetworks.TopologyConstraint(0x01, ["2a","2b"], net)
Species constraint, on tips: 2a, 2b
 stem edge number 3
 crown node number -4

julia> nni!(net , net.edge[3], true, true, [c_species2]) === nothing # we get nothing: would break species 2
true

julia> # the following gives an error, because 4,5 do not form a clade in "net"
       # PhyloNetworks.TopologyConstraint(0x02, ["4","5"], net)

julia> c_clade145 = PhyloNetworks.TopologyConstraint(0x02, ["1a","1b","1c","4","5"], net)
Clade constraint, on tips: 1a, 1b, 1c, 4, 5
 stem edge number 12
 crown node number -7

julia> nni!(net , net.edge[12], true, true, [c_species1, c_clade145]) # we get nothing (failed NNI): would break the clade
```
"""
function TopologyConstraint(type::UInt8, taxonnames::Vector{String}, net::HybridNetwork)

    ntax_clade = length(taxonnames)
    ntax_clade >= 2 ||
        error("there must be 2 or more taxon name in a clade or species constraint.")
    taxonnums = Set{Int}()
    indices = findall(in(taxonnames), [n.name for n in net.leaf])
    length(indices) == ntax_clade || error("taxa cannot be matched to nodes. Check for typos.")
    outsideclade = setdiff([n.name for n in net.leaf], taxonnames)
    # find the set of node numbers that correspond to taxonnames
    for (i,taxon) in enumerate(taxonnames)
        index = findfirst(n -> n.name == taxon, net.leaf)
        !isnothing(index) || error("taxon $taxon is not found in the network. Check for typos.")
        push!(taxonnums, net.leaf[index].number) # note: not ordered as in taxonnames
    end
    # get interior ancestor edges in major tree (no hybrids)
    matrix = hardwiredClusters(majorTree(net), vcat(taxonnames, outsideclade))
    # look for row with ones in relevant columns, zeros everywhere else (or zeros there and ones everywhere else)
    edgenum = 0 # 0 until we find the stem edge
    comparator = zeros(Int8, size(matrix)[2]-2)
    comparator[1:ntax_clade] = ones(Int8, ntax_clade)
    # go through matrix row by row to find match with comparator. Return number
    for i in 1:size(matrix, 1)
        if matrix[i,2:size(matrix)[2]-1] == comparator # found the mrca!
            edgenum = matrix[i, 1]
            break
        end
    end
    if edgenum == 0
        comparator .= 1 .- comparator
        for i in 1:size(matrix)[1]
            if matrix[i,2:size(matrix)[2]-1] == comparator
                error("""The clade given is not rooted correctly, making it a grade instead of a clade.
                You can re-try after modifying the root of your network.""")
            end
        end
        error("The taxa given do not form a clade in the network")
    end
    edgei = findfirst(e -> e.number == edgenum, net.edge)
    edgei !== nothing || error("hmm. hardwiredClusters on the major tree got an edge number not in the network")
    stemedge = net.edge[edgei]
    mrcanode = getChild(stemedge)
    TopologyConstraint(type, taxonnames, taxonnums, stemedge, mrcanode)
end
function Base.show(io::IO, obj::TopologyConstraint)
    str = (obj.type == 0x01 ? "Species constraint, on tips: " : "Clade constraint, on tips: " )
    str *= join(obj.taxonnames, ", ")
    str *= "\n stem edge number $(obj.edge.number)\n crown node number $(obj.node.number)"
    print(io, str)
end

"""
    constraintviolated(network::HybridNetwork,
                       constraints::Vector{TopologyConstraint})

True if `network` violates one (or more) of the constraints of type 1
(individuals in a species group) or type 2 (must be clades in the major tree).
Warning: constraints of type 3 are not implemented.
"""
function constraintviolated(net::HybridNetwork, constraints::Vector{TopologyConstraint})
    # fixit next PR: add option to give a vector of constraint types to check,
    #                then only check these constraint types
    if isempty(constraints)
        return false
    end # avoids extracting major tree when no constraint
    tree = majorTree(net)
    for con in constraints # checks directionality of stem edge hasn't changed
        if con.type in [0x01, 0x02]
            getChild(con.edge) === con.node || return true
            tei = findfirst(e -> e.number == con.edge.number, tree.edge)
            tei !== nothing ||
                error("hmm. edge number $(con.edge.number) was not found in the network's major tree")
            treeedge = tree.edge[tei]
            des = descendants(treeedge) # vector of node numbers, descendant tips only by default
            Set(des) == con.taxonnums || return true
        end
    end
    return false
end

"""
    updateconstraints!(constraints::Vector{TopologyConstraint}, net::HybridNetwork)

Update the set `taxonnum` in each constraint, assuming that the stem edge and
the crown node are still correct, and that their descendants are still correct.
May be needed if the node and edge numbers were modified by
[`resetNodeNumbers!`](@ref) or [`resetEdgeNumbers!`](@ref).

Warning: does *not* check that the names of leaves with numbers in `taxonnum`
are `taxonnames`.

```jldoctest
julia> net = readTopology("(((2a,2b),(((((1a,1b,1c),4),(5)#H1),(#H1,(6,7))))#H2),(#H2,10));");

julia> c_species1 = PhyloNetworks.TopologyConstraint(0x01, ["1a","1b","1c"], net)
Species constraint, on tips: 1a, 1b, 1c
 stem edge number 7
 crown node number -9

julia> c_species1.taxonnums
Set{Int64} with 3 elements:
  5
  4
  3

julia> c_clade145 = PhyloNetworks.TopologyConstraint(0x02, ["1a","1b","1c","4","5"], net)
Clade constraint, on tips: 1a, 1b, 1c, 4, 5
 stem edge number 12
 crown node number -7

julia> PhyloNetworks.resetNodeNumbers!(net)

julia> net.node[4].number = 111;

julia> PhyloNetworks.updateconstraints!([c_species1, c_clade145], net)

julia> c_species1
Species constraint, on tips: 1a, 1b, 1c
 stem edge number 7
 crown node number 21

julia> c_species1.taxonnums
Set{Int64} with 3 elements:
  5
  4
  111
```
"""
function updateconstraints!(constraints::Vector{TopologyConstraint}, net::HybridNetwork)
    if isempty(constraints)
        return nothing
    end # avoids extracting major tree when no constraint
    tree = majorTree(net)
    for con in constraints
      if con.type in [0x01, 0x02]
        getChild(con.edge) === con.node ||
            error("the stem edge and crown node have been disconnected")
        tei = findfirst(e -> e.number == con.edge.number, tree.edge)
        tei !== nothing ||
            error("stem edge number $(con.edge.number) was not found in the network's major tree")
        des = descendants(tree.edge[tei]) # vector of node numbers, descendant tips only by default
        d1 = setdiff(con.taxonnums, des)  # set
        d2 = setdiff(des, con.taxonnums)  # array
        length(d1) == length(d2) || error("missing or extra taxa in the clade / species.")
        replace!(con.taxonnums, Dict(tn => d2[i] for (i,tn) in enumerate(d1))...)
      end
    end
    return nothing
end

"""
    updateconstraintfields!(constraints::Vector{TopologyConstraint}, net::HybridNetwork)

Update fields stem `edge` and crown `node` to match the given `net`.

Assumes that the constraints are still met in `net`,
and that nodes & edges are numbered identically in `net` as in the network
used to create all `constraints`.

fixit: remove the assumption that constraints are still met, since
an NNI near the crown of a constrained clade might change the crown node
and / or the stem edge (u and v exchange).
"""
function updateconstraintfields!(constraints::Vector{TopologyConstraint}, net::HybridNetwork)
    for con in constraints
        num = con.edge.number
        con.edge = net.edge[findfirst([e.number == num for e in net.edge])]
        num = con.node.number
        con.node = net.node[findfirst([n.number == num for n in net.node])]
    end
end

#=
nice to add: option to modify branch lengths during NNI
=#

"""
    nni!(net::HybridNetwork, e::Edge, nohybridladder::Bool=true, no3cycle::Bool=true,
         constraints=TopologyConstraint[]::Vector{TopologyConstraint})

Attempt to perform a nearest neighbor interchange (NNI) around edge `e`,
randomly chosen among all possible NNIs (e.g 3, sometimes more depending on `e`)
satisfying the constraints, and such that the new network is a DAG. The number of
possible NNI moves around an edge depends on whether the edge's parent/child nodes
are tree or hybrid nodes. This is calculated by [`nnimax`](@ref).

The option `no3cycle` forbids moves that would create a 3-cycle in the network.
When `no3cycle` = false, 2-cycle and 3-cycles may be generated.

Note that the defaults values are for positional (not keyword) arguments, so
two or more arguments can be used, but in a specific order: nni!(net, e) or
nni!(net, e, nohybridladder),
nni!(net, e, nohybridladder, no3cycle),
nni!(net, e, nohybridladder, no3cycle, contraints).

Assumptions:
- The starting network does not have 3-cycles, if `no3cycle=true`.
  No check for the presence of 2- and 3-cycles in the input network.
- The edges' field `isChild1` is correct in the input network. (This field
  will be correct in the output network.)

Output:
information indicating how to undo the move or `nothing` if all NNIs failed.

# examples

```jldoctest
julia> str_network = "(((S8,S9),(((((S1,S2,S3),S4),(S5)#H1),(#H1,(S6,S7))))#H2),(#H2,S10));";

julia> net = readTopology(str_network);

julia> using Random; Random.seed!(321);

julia> undoinfo = nni!(net, net.edge[3], true, true); # true's to avoid hybrid ladders and 3-cycles

julia> writeTopology(net)
"(((S8,(((((S1,S2,S3),S4),(S5)#H1),(#H1,(S6,S7))))#H2),S9),(#H2,S10));"

julia> nni!(undoinfo...);

julia> writeTopology(net) == str_network # net back to original topology: the NNI was "undone"
true
```
"""
function nni!(net::HybridNetwork, e::Edge, nohybridladder::Bool=true, no3cycle::Bool=true,
              constraints=TopologyConstraint[]::Vector{TopologyConstraint})
    for con in constraints
        # reject NNI if the focus edge is the stem of a constraint.
        # sufficient to maintain star species constraints
        # but not sufficient for clades.
        if con.edge === e
            return nothing
        end
    end
    nnirange = 0x01:nnimax(e) # 0x01 = 1 but UInt8 instead of Int
    nnis = Random.shuffle(nnirange)
    for nummove in nnis # iterate through all possible NNIs, but in random order
        moveinfo = nni!(net, e, nummove, nohybridladder, no3cycle)
        !isnothing(moveinfo) || continue # to next possible NNI
        # if constraintviolated(net, constraints)
        #     # fixit for next PR: not needed for species constraints,
        #     # not sufficient for clades if γ / isMajor modified later, or if
        #     # constraints met after NNI but stem / crows need to be moved
        #     nni!(moveinfo...) # undo the previous NNI
        #     continue          # try again
        # end
        return moveinfo
    end
    return nothing # if we get to this point, all NNIs failed
end

"""
    nnimax(e::Edge)

Return the number of NNI moves around edge `e`,
assuming that the network is semi-directed.
Return 0 if `e` is not internal, and more generally if either node
attached to `e` does not have 3 edges.

Output: UInt8 (e.g. `0x02`)
"""
function nnimax(e::Edge)
    length(e.node[1].edge) == 3 || return 0x00
    length(e.node[2].edge) == 3 || return 0x00
    # e.hybrid and hybrid parent: RR case, 2 possible NNIs only
    # e.hybrid and tree parent:   BR case, 3 or 6 NNIs if e is may contain the root
    # e not hybrid, hyb parent:   RB case, 4 NNIs
    # e not hybrid, tree parent:  BB case, 2 NNIs if directed, 8 if undirected
    hybparent = getParent(e).hybrid
    n = (e.hybrid ? (hybparent ? 0x02 : (e.containRoot ? 0x06 : 0x03)) : # RR & BR
                    (hybparent ? 0x04 : (e.containRoot ? 0x08 : 0x02)))  # RB & BB
    return n
end

"""
    nni!(net::HybridNetwork, uv::Edge, nummove::UInt8,
         nohybridladder::Bool, no3cycle::Bool)

Modify `net` with a nearest neighbor interchange (NNI) around edge `uv`.
Return the information necessary to undo the NNI, or `nothing` if the move
was not successful (such as if the resulting graph was not acyclic (not a DAG) or if
the focus edge is adjacent to a polytomy). If the move fails, the network is not
modified.
`nummove` specifies which of the available NNIs is performed.

rooted-NNI options according to Gambette et al. (2017), fig. 8:
* BB: 2 moves, both to BB, if directed edges. 8 moves if undirected.
* RR: 2 moves, both to RR.
* BR: 3 moves, 1 RB & 2 BRs, if directed. 6 moves if e is undirected.
* RB: 4 moves, all 4 BRs.
The extra options are due to assuming a semi-directed network, whereas
Gambette et al (2017) describe options for rooted networks.
On a semi-directed network, there might be a choice of how to direct
the edges that may contain the root, e.g. choice of e=uv versus vu, and
choice of labelling adjacent nodes as α/β (BB), or as α/γ (BR).

`nohybridladder` = true prevents moves that would create a hybrid ladder in
the network, that is, 2 consecutive hybrid nodes (one parent of the other).
`no3cycle` = true prevents NNI moves that would make a 3-cycle, and assumes
that the input network does not have any 2- or 3-cycles.
If `no3cycle` is false, 3-cycles can be generated, but NNIs generating
2-cycles are prevented.

The edge field `isChild1` is assumed to be correct according to the `net.root`.
"""
function nni!(net::HybridNetwork, uv::Edge, nummove::UInt8,
              nohybridladder::Bool, no3cycle::Bool)
    nmovemax = nnimax(uv)
    if nmovemax == 0x00 # uv not internal, or polytomy e.g. species represented by polytomy with >2 indiv
        return nothing
    end
    nummove <= nmovemax || error("nummove $(nummove) must be <= $nmovemax for edge number $(uv.number)")
    u = getParent(uv)
    v = getChild(uv)
    ## TASK 1: grab the edges adjacent to uv: αu, βu, vγ, vδ and adjacent nodes
    # get edges αu & βu connected to u
    # α = u's major parent if it exists, β = u's last child or minor parent
    if u.hybrid
        αu = getMajorParentEdge(u)
        βu = getMinorParentEdge(u)
        α = getParent(αu)
        β = getParent(βu)
    else # u may not have any parent, e.g. if root node
        # pick αu = parent edge if possible, first edge of u (other than uv) otherwise
        labs = [edgerelation(e, u, uv) for e in u.edge]
        pti = findfirst(isequal(:parent), labs) # parent index: there should be 1 or 0
        ci = findall(  isequal(:child), labs)  # vector. there should be 1 or 2
        if pti === nothing
            length(ci) == 2 || error("node $(u.number) should have 2 children other than node number $(v.number)")
            pti = popfirst!(ci)
            αu = u.edge[pti]
            α = getChild(αu)
        else
            αu = u.edge[pti]
            α = getParent(αu)
        end
        βu = u.edge[ci[1]]
        β = getChild(βu)
    end
    # get edges vδ & vγ connected to v
    # δ = v's last child, γ = other child or v's parent other than u
    labs = [edgerelation(e, v, uv) for e in v.edge]
    if v.hybrid # then v must have another parent edge, other than uv
        vγ = v.edge[findfirst(isequal(:parent), labs)] # γ = getParent(vδ) ?this should be vγ right?
            #on net_hybridladder edge 1 get ERROR: ArgumentError: invalid index: nothing of type Nothing
        vδ = v.edge[findfirst(isequal(:child), labs)]
        γ = getParent(vγ)
        δ = getChild(vδ)
    else
        ci = findall(isequal(:child), labs)
        length(ci) == 2 || error("node $(v.number) should have 2 children")
        vγ = v.edge[ci[1]] # γ = getChild(vγ)
        vδ = v.edge[ci[2]]
        γ = getChild(vγ)
        δ = getChild(vδ)
    end
    ## TASK 2: semi-directed network swaps:
    ## swap u <-> v, α <-> β if undirected and according to move number
    if !u.hybrid && uv.containRoot # case BB or BR, with uv that may contain the root
        nmoves = 0x03 # number of moves in the rooted (directed) case
        if !v.hybrid # BB, uv and all adjacent edges are undirected: 8 moves
            if nummove > 0x04  # switch u & v with prob 1/2
                nummove -= 0x04  # now in 1,2,3,4
                (u,v) = (v,u)
                (α,αu,β,βu,γ,vγ,δ,vδ) = (γ,vγ,δ,vδ,α,αu,β,βu)
            end
            nmoves = 0x02
        end
        # now nmoves should be in 1, ..., 2*nmoves: 1:4 (BB) or 1:6 (BR) #? should this be 2*nummoves?
        # next: switch α & β with probability 1/2
        if nummove > nmoves
            nummove -= nmoves # now nummoves in 1,2 (BB) or 1,2,3 (BR)
            (α,αu,β,βu) = (β,βu,α,αu)
        end
    end
    ## TASK 3: rooted network swaps, then do the NNI
    if u.hybrid # RR or RB cases
        # swap α & β: reduce NNIs from 2 to 1 (RR) or from 4 to 2 (RB)
        if nummove == 0x02 || nummove == 0x04
            (α,αu,β,βu) = (β,βu,α,αu)
        end
        if !v.hybrid && nummove > 0x02 # case RB, moves 3 or 4
            (γ,vγ,vδ,δ) = (δ,vδ,vγ,γ)
        end
        # detach β and graft onto vδ
        if no3cycle # don't create a 3-cycle
            if problem4cycle(α,γ, β,δ) return nothing; end
        else # don't create a 2-cycle
            if α === γ || β === δ return nothing; end
        end
        res = nni!(αu, u, uv, v, vδ)
    else # u = bifurcation: BB or BR cases
        if !v.hybrid # case BB: 2 options
            if nummove == 0x02
                (γ,vγ,vδ,δ) = (δ,vδ,vγ,γ)
            end
            # detach β and graft onto vδ
            if no3cycle
                if problem4cycle(α,γ, β,δ) return nothing; end
            elseif α === γ || β === δ return nothing
            end
            res = nni!(αu, u, uv, v, vδ)
        else # case BR: 3 options
            # DAG check
            # if α parent of u, check if β -> γ (directed or undirected)
            # if undirected with β parent of u, check α -> γ
            # if undirected and u = root:
            #   moves 1 and 3 will fail if α -> γ
            #   moves 2 and 3 will fail if β -> γ
            # nummove 3 always creates a nonDAG when u is the root
            αparentu = getChild(αu)===u
            βparentu = getChild(βu)===u
            if αparentu
                if γ === β || isdescendant(γ, β) return nothing; end
            elseif βparentu
                if γ === α || isdescendant(γ, α) return nothing; end
            else # cases when u is root
                if nummove == 0x01 && (γ === α || isdescendant(γ, α))
                    return nothing
                elseif nummove == 0x02 && (γ === β || isdescendant(γ, β))
                    return nothing
                elseif nummove == 0x03 # fail: not DAG
                    return nothing
                end
            end
            if nummove == 0x01 # graft γ onto α
                if no3cycle
                    if problem4cycle(α,γ, β,δ) return nothing; end
                elseif α === γ || β === δ return nothing
                end
                if nohybridladder && α.hybrid return nothing; end
                res = nni!(vδ, v, uv, u, αu)
            elseif nummove == 0x02 # graft γ onto β
                if nohybridladder && β.hybrid return nothing; end
                if no3cycle
                    if problem4cycle(β,γ, α,δ) return nothing; end
                elseif β === γ || α === δ return nothing
                end
                res = nni!(vδ, v, uv, u, βu)
            else # nummove == 0x03
                if αparentu # if α->u: graft δ onto α
                    if no3cycle
                        if problem4cycle(α,δ, β,γ) return nothing; end
                    elseif β === γ || α === δ return nothing
                    end
                    if nohybridladder && α.hybrid return nothing; end
                    res = nni!(vγ, v, uv, u, αu)
                else # if β->u: graft δ onto β
                    if no3cycle
                        if problem4cycle(α,γ, β,δ) return nothing; end
                    elseif α === γ || β === δ return nothing
                    end
                    # we could check nohybridladder && β.hybrid but no need because
                    # we only get here if uv.containRoot so β cannot be hybrid
                    res = nni!(vγ, v, uv, u, βu)
                # case when u is root already excluded (not DAG)
                end
            end
        end
    end
    return res
end

"""
    nni!(αu, u, uv::Edge, v, vδ)
    nni!(αu,u,uv,v,vδ, flip::Bool, inner::Bool, indices)

Transform a network locally around the focus edge `uv` with
the following NNI, that detaches u-β and grafts it onto vδ:

    α - u -- v ------ δ
        |    |
        β    γ

    α ------ v -- u - δ
             |    |
             γ    β

`flip` boolean indicates if the uv edge was flipped
`inner` boolean indicates if edges αu and uv both point toward node u,
i.e. α->u<-v<-δ. If this is true, we flip the hybrid status of αu and vδ.

`indices` give indices for nodes and edges u_in_αu, αu_in_u, vδ_in_v, and v_in_vδ.
These are interpreted as:

    u_in_αu: the index for u in the edge αu
    αu_in_u: the index for αu in node u
    vδ_in_v: the index for vδ in node v
    v_in_vδ: the index for v in edge vδ

**Warnings**:

- *No* check of assumed adjacencies
- Not implemented for cases that are not necessary thanks to symmetry,
  such as cases covered by `nni!(vδ, v, uv, u, αu)` or `nni!(βu, u, v, vγ)`.
  More specifically, these cases are not implemented (and not checked):
  * u not hybrid & v hybrid
  * u hybrid, v not hybrid, α -> u <- v -> δ
- Because of this, `nni(αu,u,uv,v,vδ, ...)` should not be used directly;
  use instead `nni!(net, uv, move_number)`.
- nni!(undoinfo...) restores the topology, but edges below hybrid nodes
  will have length 0.0 even if they didn't before.

Node numbers and edge numbers are not modified.
Edge `uv` keeps its direction unchanged *unless* the directions were
`α -> u -> v -> δ` or `α <- u <- v <- δ`,
in which case the direction of `uv` is flipped.

The second version's input has the same signature as the output, but
will undo the NNI more easily. This means that if `output = nni!(input)`,
then `nni!(output...)` is valid and undoes the first operation.

Right now, branch lengths are not modified except when below a hybrid node.
Future versions might implement options to modify branch lengths.
"""
function nni!(αu::Edge, u::Node, uv::Edge, v::Node, vδ::Edge)
    # find indices to detach αu from u, and to detach v from vδ
    u_in_αu = findfirst(n->n===u, αu.node)
    αu_in_u = findfirst(e->e===αu, u.edge)
    vδ_in_v = findfirst(e->e===vδ, v.edge)
    v_in_vδ = findfirst(n->n===v, vδ.node)
    # none of them should be 'nothing' --not checked
    αu_child = getChild(αu)
    uv_child = getChild(uv)
    vδ_child = getChild(vδ)
    # flip the direction of uv if α->u->v->δ or α<-u<-v<-δ
    flip = (αu_child === u && uv_child === v && vδ_child !== v) ||
             (αu_child !== u && uv_child === u && vδ_child === v)
    inner = (flip ? false : αu_child === u && vδ_child === v )
    # fixit: calculate new edge lengths here; give them as extra arguments to nni! below
    res = nni!(αu,u,uv,v,vδ, flip,inner, u_in_αu,αu_in_u,vδ_in_v,v_in_vδ)
    return res
end
function nni!(αu::Edge, u::Node, uv::Edge, v::Node, vδ::Edge,
              flip::Bool, inner::Bool,
              u_in_αu::Int, αu_in_u::Int, vδ_in_v::Int, v_in_vδ::Int)
              # extra arguments: new edge lengths
    # fixit: extract old edge lengths; update to new lengths; return old lengths as extra items
    ## TASK 1: do the NNI swap: to detach & re-attach at once
    (αu.node[u_in_αu], u.edge[αu_in_u], v.edge[vδ_in_v], vδ.node[v_in_vδ]) =
        (v, vδ, αu, u)
    ## TASK 2: update edges' directions
    if flip
        uv.isChild1 = !uv.isChild1
    end
    ## TASK 3: update hybrid status of nodes (#? is it ever nodes?)
             # & edges, and edges containRoot field
    if u.hybrid && !v.hybrid
        if flip # RB, BR 1 or BR 2': flip hybrid status of uv and (αu or vδ)
            if uv.hybrid # BR 1 or BR 2'
                vδ.hybrid = true
                vδ.gamma = uv.gamma
                vδ.isMajor = uv.isMajor
                # length switches: for keeping the network unzipped. could be option later
                αu.length = uv.length
                uv.length = 0.0
                uv.hybrid = false
                uv.gamma = 1.0
                uv.isMajor = true
                norootbelow!(uv)
            else # RB
                uv.hybrid = true
                uv.gamma = αu.gamma
                uv.isMajor = αu.isMajor
                # length switches: for keeping the network unzipped. could be option later
                uv.length = vδ.length
                vδ.length = 0.0
                αu.hybrid = false
                αu.gamma = 1.0
                αu.isMajor = true
                if αu.containRoot
                    allowrootbelow!(v, αu)
                end
            end
        else # assumes α<-u or v<-δ --- but not checked
            # not implemented: α->u<-v->δ: do uv.hybrid=false and γv.hybrid=true
            if inner # BR 3 or 4': α->u<-v<-δ: flip hybrid status of αu and vδ
                # no lengths to switch in these cases
                vδ.hybrid = true
                vδ.gamma = αu.gamma
                vδ.isMajor = αu.isMajor
                αu.hybrid = false
                αu.gamma = 1.0
                αu.isMajor = true
                # containRoot: nothing to update
            else # BR 1' or 2: α<-u<-v->δ : hybrid status just fine
                # length switches: for keeping the network unzipped. could be option later
                αu.length = vδ.length
                vδ.length = 0.0
                norootbelow!(vδ)
                if uv.containRoot
                    allowrootbelow!(αu)
                end
            end
        end
    elseif u.hybrid && v.hybrid # RR: hybrid edges remain hybrids, but switch γs
        # we could have -αu-> -uv-> -vδ-> or <-αu- <-uv- <-vδ-
        hyb = ( getChild(αu) == v ? αu : vδ ) # uv's hybrid parent edge: either αu or vδ
        (uv.gamma,   hyb.gamma)   = (hyb.gamma,   uv.gamma)
        (uv.isMajor, hyb.isMajor) = (hyb.isMajor, uv.isMajor)
    end
    # throw error if u not hybrid & v hybrid? (does not need to be implemented)
    return vδ, u, uv, v, αu, flip, inner, v_in_vδ, αu_in_u, vδ_in_v, u_in_αu
end

"""
    problem4cycle(β::Node, δ::Node, α::Node, γ::Node)

Check if the focus edge uv has a 4 cycle that could lead to a 3 cycle
after an chosen NNI. Return true if there is a problem 4 cycle, false if none.
"""
function problem4cycle(β::Node, δ::Node, α::Node, γ::Node)
    isconnected(β, δ) || isconnected(α, γ)
end

"""
    checkspeciesnetwork!(network::HybridNetwork,
                         constraints::Vector{TopologyConstraint})

Check that the network satisfies a number of requirements:
- no polytomies, other than at species constraints: throws an error otherwise
- no unnecessary nodes: fuse edges at nodes with degree two, including at the root
  (hence the bang: `net` may be modified)
- topology constraints are met.

Output: true if all is good, false if one or more clades are violated.
"""
function checkspeciesnetwork!(net::HybridNetwork, constraints::Vector{TopologyConstraint})
    for n in net.node # check for no polytomies: all nodes should have up to 3 edges
        if length(n.edge) > 3 # polytomies allowed at species constraints only
            coni = findfirst(c -> c.type == 1 && c.node === n, constraints)
            coni !== nothing ||
                error("The network has a polytomy at node number $(n.number). Please resolve.")
        end
    end
    removedegree2nodes!(net)
    return !constraintviolated(net, constraints)
end

"""
    mapindividuals(net::HybridNetwork, mappingFile::String)

Return a network expanded from `net`, where species listed in the mapping file
are replaced by individuals mapped to that species.
If a species has only 1 individual, the name of the leaf for that species is
replaced by the name of its one individual representative.
If a species has 2 or more individuals, the leaf for that species is
expanded into a "star" (polytomy if 3 or more individuals) with a tip for each
individual. If a species is in the network but not listed in the mapping file,
the tip for that species is left as is. Species listed in the mapping file
but not present in the network are ignored.

The mapping file should be readable by `CSV.File` and contain two columns:
one for the species names and one for the individual (or allele) names.
fixit: make this function more flexible by accepting column names

Output: individual-level network and vector of species constraint(s).

# examples

```jldoctest
julia> species_net = readTopology("(((S8,S9),((((S1,S4),(S5)#H1),(#H1,(S6,S7))))#H2),(#H2,S10));");

julia> filename = joinpath(dirname(Base.find_package("PhyloNetworks")), "..", "examples", "mappingIndividuals.csv");

julia> filename |> read |> String |> print # to see what the mapping file contains
species,individual
S1,S1A
S1,S1B
S1,S1C

julia> individual_net, species_constraints = mapindividuals(species_net, filename);

julia> writeTopology(individual_net, internallabel=true)
"(((S8,S9),(((((S1A,S1B,S1C)S1,S4),(S5)#H1),(#H1,(S6,S7))))#H2),(#H2,S10));"

julia> species_constraints
1-element Vector{PhyloNetworks.TopologyConstraint}:
 Species constraint, on tips: S1A, S1B, S1C
 stem edge number 4
 crown node number 3
```
"""
function mapindividuals(net::HybridNetwork, mappingFile::String)
    mappingDF = DataFrame(CSV.File(mappingFile); copycols=false)
    specieslist = unique(mappingDF[:, 1])
    individualnet = deepcopy(net)
    constraints = TopologyConstraint[]
    for species in specieslist
        individuals = mappingDF[mappingDF.species .== species, 2]
        res = addindividuals!(individualnet, species, individuals)
        res === nothing || # nothing if species not found in the network, or 0,1 individuals
            push!(constraints, res)
    end
    return individualnet, constraints
end

"""
    addindividuals!(net::HybridNetwork, species::AbstractString, individuals::Vector)

Add individuals to their species leaf as a star, and return the
corresponding species constraint. `nothing` is returned in 2 cases:

- if `individuals` contains only 1 individual, in which case the
  name of the leaf for that species is replaced by the name of
  its one individual representative
- if `species` is not found in the network

Spaces in individuals' names are eliminated, see [`cleantaxonname`](@ref).

Called by [`mapindividuals`](@ref).

# examples

```jldoctest
julia> net = readTopology("(S8,(((S1,S4),(S5)#H1),(#H1,S6)));");

julia> PhyloNetworks.addindividuals!(net, "S1", ["S1A", "S1B", "S1C"])
Species constraint, on tips: S1A, S1B, S1C
 stem edge number 2
 crown node number 2

julia> writeTopology(net) # 3 new nodes, S1 now internal: not a tip
"(S8,((((S1A,S1B,S1C)S1,S4),(S5)#H1),(#H1,S6)));"
```
"""
function addindividuals!(net::HybridNetwork, species::AbstractString, individuals::Vector{String})
    length(individuals) > 0 || return nothing
    species = strip(species)
    speciesnodeindex = findfirst(l -> l.name == species, net.node)
    if speciesnodeindex === nothing
        @warn "species $species not found in network (typo in mapping file?). will be ignored."
        return nothing
    end
    if length(individuals) == 1 # change species name into name of individual representative
        net.node[speciesnodeindex].name = cleantaxonname(individuals[1])
        return nothing
    end
    # there are 2+ individuals: replace the species tip by a star, will need a species constraint
    taxnames = map(cleantaxonname, individuals)
    for tax in taxnames
        addleaf!(net, net.node[speciesnodeindex], tax);
    end
    return TopologyConstraint(0x01, taxnames, net)
end

"""
    cleantaxonname(taxonname::AbstractString)

Return a String with leading and trailing spaces removed, and interior spaces
replaced by underscores: good for using as tip names in a network without
causing future error when reading the newick description of the network.
"""
function cleantaxonname(taxonname::AbstractString)
    tname = string(strip(taxonname)) # SubString if we don't do string()
    m = match(r"\s", tname)
    if m !== nothing
        @warn """Spaces in "$tname" may cause errors in future network readability: replaced by _"""
        tname = replace(tname, r"\s+" => "_")
    end
    return tname
end

"""
    moveroot!(net::HybridNetwork, constraints=TopologyConstraint[]::Vector{TopologyConstraint})

Move the root to a randomly chosen non-leaf node that is different from
the current root, and not within a constraint clade or species.
Output: `true` if successul, `nothing` otherwise.
"""
function moveroot!(net::HybridNetwork, constraints=TopologyConstraint[]::Vector{TopologyConstraint})
    newrootrandomorder = Random.shuffle(1:length(net.node))
    oldroot = net.root
    for newrooti in newrootrandomorder
        newrooti != oldroot || continue
        newrootnode = net.node[newrooti]
        newrootnode.leaf && continue # try next potential new root if current is a leaf
        # Check the new root is NOT at the top or inside contraint clade or within pecies
        newrootfound = true
        for con in constraints
            if con.node === newrootnode || isdescendant(newrootnode, con.node) # assumes the constraint is met
                newrootfound = false
                break # out of constraint loop
            end
        end
        newrootfound || continue # to next potential new root
        # Check for no directional conflict
        try
            net.root = newrooti
            directEdges!(net)
            # warning: assumes that the previous root has degree 3
            # otherwise, need to delete previous root by fusing its 2 adjacent edges
        catch e # RootMismatch error if root placement is below a hybrid
            isa(e, RootMismatch) || rethrow(e)
            net.root = oldroot
            directEdges!(net) # revert edges' directions to match original rooting
            continue # to next potential new root
        end
        # if we get here: newrooti passed all the checks
        return true
    end
    # if we get here: none of the root positions worked
    return nothing
end

"""
    fliphybrid!(net::HybridNetwork, minor=true::Bool, nohybridladder=false::Bool,
                constraints=TopologyConstraint[]::Vector{TopologyConstraint})

Cycle through hybrid nodes in random order until an admissible flip is found.
At this hybrid node, flip the indicated hybrid parent edge (minor or major).

If an admissible flip is found, return the tuple: newhybridnode, flippededge, oldchildedge.
Otherwise, return nothing.

The flip can be undone with
`fliphybrid!(net, newhybridnode, minor, constraints)`, or
`fliphybrid!(net, newhybridnode, !flippededge.isMajor, constraints)`
more generally, such as if the flipped edge had its γ modified after the
original flip.

*Warnings*
- if the root needed to be reset and if the original root was of degree 2,
  then a node of degree 2 remains in the modified network.
- undoing the flip may not recover the original root in case
  the root position was modified during the original flip.
"""
function fliphybrid!(net::HybridNetwork, minor=true::Bool,
                     nohybridladder=false::Bool,
                     constraints=TopologyConstraint[]::Vector{TopologyConstraint})
    hybridindex = Random.shuffle(1:length(net.hybrid)) # indices in net.hybrid
    while !isempty(hybridindex) # all minor edges
        i = pop!(hybridindex)
        undoinfo = fliphybrid!(net, net.hybrid[i], minor, nohybridladder, constraints)
        if !isnothing(undoinfo) # if it failed, the network was not changed
            return undoinfo
        end # else, continue until explored all i in hybridindex
    end
    return nothing
end

"""
    fliphybrid!(net::HybridNetwork, hybridnode::Node, minor=true::Bool,
                nohybridladder=false::Bool,
                constraints=TopologyConstraint[]::Vector{TopologyConstraint})

Flip the direction of a single hybrid edge:
the minor parent edge of `hybridnode` by default,
or the major parent edge if `minor` is false.
The parent node of the hybrid edge becomes the new hybrid node.
The former hybrid edge partner is converted to a tree edge (with γ=1),
and `hybridnode` becomes a tree node.

For the flip to be admissible, the new network must be a semi-directed
phylogenetic network: with a root such that the rooted version is a DAG.
If `nohybridladder` is false (default), the flip may create a hybrid ladder
If `nohybridladder` is true and if the flip would create a hybrid ladder,
then the flip is not admissible.
A hybrid ladder is when a hybrid child of another hybrid.

The new hybrid partner is an edge adjacent to the new hybrid node,
such that the flip is admissible (so it must be a tree edge).
The flipped edge retains its original γ.
The new hybrid edge is assigned inheritance 1-γ.

Output:
`(newhybridnode, flippededge, oldchildedge)` if the flip is admissible,
`nothing` otherwise.

The network is unchanged if the flip is not admissible.
If the flip is admissible, the root position may be modified, and
the direction of tree edges (via `isChild1`) is modified accordingly. If the
root needs to be modified, then the new root is set to the old hybrid node.

The index of the new hybrid node in `net.hybrid` is equal to that of the
old `hybridnode`.

Warning: Undoing this move may not recover the original root if
the root position was modified.
"""
function fliphybrid!(net::HybridNetwork, hybridnode::Node, minor=true::Bool,
                     nohybridladder=false::Bool,
                     constraints=TopologyConstraint[]::Vector{TopologyConstraint})
    #= for species constraints, there is nothing to check, because hybrids cannot point into or come out of the group
    for con in constraints
        if con.type == # 0x02/0x03 for types 2 and 3, need to consider more cases
            return nothing
        end
    end
    =#
    runDirectEdges = false
    edgetoflip, edgetokeep = minor ?
        (getMinorParentEdge(hybridnode), getMajorParentEdge(hybridnode)) :
        (getMajorParentEdge(hybridnode), getMinorParentEdge(hybridnode))
    oldchildedge = getChildEdge(hybridnode)
    newhybridnode = getParent(edgetoflip)
    if newhybridnode.hybrid # already has 2 parents: cannot had a third.
        return nothing
    end
    ## choose newhybridedge ##
    p2 = getParent(edgetokeep) # parent node
    isdesc = Bool[]
    for e in newhybridnode.edge # is p2 undirected descendant of nhn via this edge?
        isp2desc = false
        #= isp2desc = true means:
        there is a semi-directed path  newhybridnode -e-> neighbor --...-> p2
        so: flipping hybridedge without making e the new partner would:
        - make e a child of newhybridnode, and
        - create a directed cycle: p2 -hybridedge-> newhybridnode -e-> neibr --...-> p2
        so we HAVE to make e the new partner hybrid edge if its "isp2desc" is true
        =#
        if e !== edgetoflip # e cannot be hybrid --e-> nhn because earlier check
            neibr = getOtherNode(e, newhybridnode) # neighbor of new hybrid node via e
            if !neibr.leaf && (p2 === neibr ||
                isdescendant_undirected(p2, neibr, e))
                isp2desc = true
            end
        end
        push!(isdesc, isp2desc)
    end
    sum_isdesc = sum(isdesc)
    if sum_isdesc > 1 # if 2+ edges have a semi-directed path to p2,
        # flipping either would create a semi-directed cycle via the other(s)
        return nothing
    end
    if sum_isdesc == 0 # no risk to create directed cycle.
        # choose the current parent edge of nhn to: keep its current direction
        # and keep reaching all nodes from a single root
        newhybridedge = getMajorParentEdge(newhybridnode)
    else
        newhybridedge = newhybridnode.edge[findfirst(isdesc)]
        if newhybridedge.hybrid # cannot flip another hybrid edge, but we needed
            return nothing      # to calculate its isp2desc for total # of true's
        end
    end
    # @debug "edgetoflip is $edgetoflip, edgetokeep is $edgetokeep, newhybridedge is $newhybridedge"
    if nohybridladder
        # case when edgetoflip would be bottom rung of ladder
        if getOtherNode(newhybridedge, newhybridnode).hybrid
            return nothing
        end
        # case when edgetoflip would be the top rung: happens when
        # newhybridnode = center point of a W structure initially
        for e in newhybridnode.edge
            if e !== edgetoflip && e.hybrid # newhybridedge is NOT hybrid, so far
                return nothing
            end
        end
    end
    # change hybrid status and major status for nodes and edges
    hybridnode.hybrid = false
    edgetokeep.hybrid = false
    newhybridnode.hybrid = true
    newhybridedge.hybrid = true
    newhybridedge.isMajor = minor
    edgetokeep.isMajor = true
    # update node order to keep isChild1 attribute of hybrid edges true
    edgetoflip.isChild1 = !edgetoflip.isChild1 # just switch
    if getChild(newhybridedge) !== newhybridnode # includes the case when the newhybridnode was the root
        # then flip newhybridedge too: make it point towards newhybridnode
        newhybridedge.isChild1 = !newhybridedge.isChild1
        net.root = findfirst(n -> n === hybridnode, net.node)
        runDirectEdges = true # many edges are likely to need to have directions flipped here
    end
    # give new hybridnode a name
    newhybridnode.name = hybridnode.name
    hybridnode.name = "" # remove name from former hybrid node
    # update gammas
    newhybridedge.gamma = edgetokeep.gamma
    edgetokeep.gamma = 1.0
    # update hybrids in network (before directEdges!)
    hybridindex = findfirst(n -> n === hybridnode, net.hybrid)
    net.hybrid[hybridindex] = newhybridnode
    if runDirectEdges # direct edges only if root is moved
        directEdges!(net)
    else # if not, only update containRoot attributes
        # norootbelow in child edges of newhybridedge
        for ce in newhybridnode.edge
            ce !== newhybridedge || continue # skip e
            getParent(ce) === newhybridnode || continue # skip edges that aren't children of cn
            norootbelow!(ce)
        end
        # allow root for edges below old hybrid node
        if edgetokeep.containRoot #else, already forbidden below due to hybrid above
            allowrootbelow!(hybridnode, edgetokeep) # old hybrid node
        end
    end
    return newhybridnode, edgetoflip, oldchildedge
end


text:
mutable struct Poly
    T::Type
    x
    n::Int
    a₀
    terms::Union{Dict{Int,Any},Nothing}

    Poly(val) = new(typeof(val), nothing, 0, val, nothing)
    Poly(T::Type, x) = new(T, value(x), 0, 0, nothing)
    Poly(T::Type, x, n, a₀, terms) = new(T, value(x), n, a₀, terms)
end

# a series of helper functions to facilitate generation of low-degree polynomials
function gen_poly(T::Type, x, a...)
    p = Poly(T, x)

    for (i, val) in enumerate(a)
        p[i-1] = val
    end

    update_order!(p)
    p
end

gen_poly(x, a...) = gen_poly(Any, x, a...)

function degree(p::Poly)
    d = p.n
    while d > 0 && iszero(p[d])
        d -= 1
    end
    d
end

degree(p::Poly, x) = isequal(p.x, x) ? degree(p) : degree(sym(p), x)

"""
    poly_degree returns the degree of a polynomial equation eq
"""
function degree(poly, x)
    eqs = collect_powers(poly, x)
    round(Int, maximum(keys(eqs)))
end

terms(p::Poly) = p.terms
var(p::Poly) = p.x
leading(p::Poly) = degree(p) > 0 ? p[degree(p)] : p.a₀
cont(p::Poly) = (p.n > 0 ? gcd(p.a₀, values(p.terms)...) : p.a₀) * sign(leading(p))
prim(p::Poly) = p / cont(p)

function update_order!(p::Poly)
    p.n == 0 && return 0
    while iszero(p[p.n]) && p.n>0
        delete!(p.terms, p.n)
        p.n -= 1
    end
    if p.n == 0
        p.terms = nothing
    end
    p
end

function poly(T::Type, eq, x)
    x === nothing && return Poly(T, nothing, 0, eq, nothing)
    terms = collect_powers(eq, x)
    terms = Dict{Int,Any}(convert(Int,k) => val for (k,val) in terms)

    if haskey(terms, 0)
        a₀ = terms[0]
        delete!(terms, 0)
    else
        a₀ = 0
    end
    n = isempty(terms) ? 0 : maximum(keys(terms))
    Poly(T, x, n, a₀, terms)
end

function poly(T::Type, eq)
    xs = get_variables(eq)

    if length(xs) > 1
        error("more than one implicit variable. Pass x to polymerize.")
    elseif length(xs) == 0
        return nothing
    else
        return poly(T, eq, xs[1])
    end
end

poly(eq, x) = poly(Any, eq, x)
poly(eq) = poly(Any, eq)
poly(T::Type, p::Poly) = convert(T, p)
poly(p::Poly) = p

polyas(p::Poly, val) = poly(p.T, val, p.x)

function Base.getindex(p::Poly, k::Integer)
    if k == 0
        return convert(p.T, p.a₀)
    elseif k > p.n || !haskey(p.terms, k)
        return convert(p.T, 0)
    else
        return convert(p.T, p.terms[k])
    end
end

function Base.setindex!(p::Poly, term, k::Number)
    if k == 0
        p.a₀ = term
    elseif p.terms !== nothing && iszero(term)
        delete!(p.terms, k)
        if k == p.n
            update_order!(p)
        end
    else
        if p.terms === nothing
            p.terms = Dict{Int,Any}()
        end
        p.terms[k] = term
        p.n = max(p.n, k)
    end
end

simplify_rational(x::Rational) = denominator(x) == 1 ? numerator(x) : x
simplify_rational(x) = x

sym(p::Poly) = (p.x === nothing ? p.a₀ : sum([simplify_rational(p[k])*p.x^k for k=0:p.n]))
sym(x) = x

function Base.iszero(p::Poly)
    !iszero(simplify(p.a₀)) && return false
    p.n > 0 && return all(iszero, values(p.terms))
    true
end

Base.show(io::IO, p::Poly) = print(io, sym(p))
Base.copy(p::Poly) = Poly(p.T, p.x, p.n, p.a₀, p.n>0 ? copy(p.terms) : nothing)

function extract_coef(p::Poly)
    c = p.T[p.a₀]
    if p.n > 0
        append!(c, values(p.terms))
    end
    c
end

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

# pox (power-of-x) is a symbolic function to keep track of the powers of x
# pox(k,n) means k*x^n
@syms pox(k, n)

is_pox(x) = istree(x) && operation(x)==pox
is_not_pox(x) = !is_pox(x)

get_coef(p) = is_pox(p) ? arguments(p)[1] : p
get_power(p) = is_pox(p) ? arguments(p)[2] : 0

replace_x(eq, x) = substitute(eq, Dict(x => pox(1,1)))

iscomplex(x) = x isa Complex

count_rule1 = @rule ^(pox(~k, ~n1), ~n2) => isequal(~k,1) ? pox(1, ~n1 * ~n2) : pox(^(~k,~n2), ~n1 * ~n2)
count_rule2 = @rule pox(~k1, ~n1) * pox(~k2, ~n2) => pox(~k1 * ~k2, ~n1 + ~n2)
count_rule3 = @acrule pox(~k, ~n) * ~u::is_not_pox => pox(~k * ~u, ~n)

"""
    collect_powers separates the powers of x in eq (a polynomial) and returns
    a dictionary of power => term
"""
function collect_powers(eq, x)
    eq = expand(expand_derivatives(eq))
    eq = replace_x(eq, x)
    #eq = Prewalk(PassThrough(count_rule1))(eq)
    eq = Fixpoint(Prewalk(PassThrough(Chain([count_rule1, count_rule2, count_rule3]))))(eq)

    if !istree(eq)
        return Dict{Any, Any}(0 => eq)
    elseif is_pox(eq)
        return Dict{Any, Any}(get_power(eq) => get_coef(eq))
    else
        eqs = Dict{Any, Any}()
        for term in arguments(eq)
            n = get_power(term)
            if haskey(eqs, n)
                eqs[n] = eqs[n] + get_coef(term)
            else
                eqs[n] = get_coef(term)
            end
        end

        return eqs
    end
end


text:
# This file is a part of Julia. License is MIT: http://julialang.org/license

# This script does the following
# Sets up n unix pipes
# For the odd pipes, a byte is written to the write end at intervals specified in intvls
# Nothing is written into the even numbered pipes
# Odd numbered pipes are tested for reads
# Even numbered pipes are tested for timeouts
# Writable ends are always tested for writability before a write

n = 20
intvls = [2, .2, .1, .005]

pipe_fds = Vector{Any}(n)
for i in 1:n
    @static if is_windows()
        pipe_fds[i] = Array{Libc.WindowsRawSocket}(2)
        0 == ccall(:wsasocketpair, Cint, (Cint, Cuint, Cint, Ptr{Libc.WindowsRawSocket}), 1, 1, 6, pipe_fds[i]) || error(Libc.FormatMessage())
    else
        pipe_fds[i] = Array{RawFD}(2)
        @test 0 == ccall(:pipe, Cint, (Ptr{RawFD},), pipe_fds[i])
    end
end

function pfd_tst_reads(idx, intvl)
    global ready += 1
    wait(ready_c)
    tic()
    start_evt2 = Condition()
    evt2 = @async (notify(start_evt2); poll_fd(pipe_fds[idx][1], intvl; readable=true, writable=false))
    wait(start_evt2); yield() # make sure the async poll_fd is pumping events
    evt = poll_fd(pipe_fds[idx][1], intvl; readable=true, writable=false)
    t_elapsed = toq()
    @test !evt.timedout
    @test evt.readable
    @test !evt.writable
    @test evt === wait(evt2)

    # println("Expected ", intvl, ", actual ", t_elapsed, ", diff ", t_elapsed - intvl)
    # Disabled since this assertion fails randomly, notably on build VMs (issue #12824)
    # @test t_elapsed <= (intvl + 1)

    dout = Array{UInt8}(1)
    @static if is_windows()
        1 == ccall(:recv, stdcall, Cint, (Ptr{Void}, Ptr{UInt8}, Cint, Cint), pipe_fds[idx][1], dout, 1, 0) || error(Libc.FormatMessage())
    else
        @test 1 == ccall(:read, Csize_t, (Cint, Ptr{UInt8}, Csize_t), pipe_fds[idx][1], dout, 1)
    end
    @test dout[1] == Int8('A')
end

function pfd_tst_timeout(idx, intvl)
    global ready += 1
    wait(ready_c)
    tic()
    start_evt2 = Condition()
    evt2 = @async (notify(start_evt2); poll_fd(pipe_fds[idx][1], intvl; readable=true, writable=false))
    wait(start_evt2); yield() # make sure the async poll_fd is pumping events
    evt = poll_fd(pipe_fds[idx][1], intvl; readable=true, writable=false)
    @test evt.timedout
    @test !evt.readable
    @test !evt.writable
    @test evt === wait(evt2)
    t_elapsed = toq()

    # Disabled since these assertions fail randomly, notably on build VMs (issue #12824)
    # @test intvl <= t_elapsed
    # @test t_elapsed <= (intvl + 1)
end

# Odd numbers trigger reads, even numbers timeout
for (i, intvl) in enumerate(intvls)
    @sync begin
        global ready = 0
        global ready_c = Condition()
        t = Vector{Task}(n)
        for idx in 1:n
            if isodd(idx)
                t[idx] = @async pfd_tst_reads(idx, intvl)
            else
                t[idx] = @async pfd_tst_timeout(idx, intvl)
            end
        end

        while ready < n
            sleep(0.1)
        end
        ready = 0
        # tickle only the odd ones, but test for writablity for everyone
        for idx in 1:n
            event = poll_fd(pipe_fds[idx][2], 0.001; readable=true, writable=true)
            @test !event.timedout
            @test !event.readable
            @test event.writable

            if isodd(idx)
                @static if is_windows()
                    1 == ccall(:send, stdcall, Cint, (Ptr{Void}, Ptr{UInt8}, Cint, Cint), pipe_fds[idx][2], "A", 1, 0) || error(Libc.FormatMessage())
                else
                    @test 1 == ccall(:write, Csize_t, (Cint, Ptr{UInt8}, Csize_t), pipe_fds[idx][2], "A", 1)
                end
            end
        end
        notify(ready_c, all=true)
        for idx in 1:n
            wait(t[idx])
        end
    end
end

for i in 1:n
    for j = 1:2
        @static if is_windows()
            0 == ccall(:closesocket, stdcall, Cint, (Ptr{Void},), pipe_fds[i][j]) || error(Libc.FormatMessage())
        else
            @test 0 == ccall(:close, Cint, (Cint,), pipe_fds[i][j])
        end
    end
end

# issue #12473
# make sure 1-shot timers work
let a = []
    Timer(t->push!(a, 1), 0.01, 0)
    sleep(0.2)
    @test a == [1]
end


text:
gaussian_2d(x = -3:0.01:3, y = -7:0.01:3; x₀ = 0, y₀ = -2, σx = 1, σy = 2) = begin
    X = repeat(x', length(y), 1)
    Y = repeat(y, 1, length(x))
    x, y, map((x, y) -> exp(-((x - x₀) / 2σx)^2 - ((y - y₀) / 2σy)^2), X, Y)
end

@testset "arbitrary colormap" begin
    colormap = :cividis
    p = @inferred contourplot(gaussian_2d()...; colormap = colormap)
    test_ref("references/contourplot/gauss_$colormap.txt", @show_col(p))
end

@testset "number of levels" begin
    levels = 5
    p = @inferred contourplot(gaussian_2d()...; levels = levels)
    test_ref("references/contourplot/gauss_$(levels)levels.txt", @show_col(p))
end

@testset "update contourplot" begin
    p = @inferred contourplot(gaussian_2d()...; levels = 2)
    contourplot!(p, gaussian_2d(; σx = 0.5, σy = 0.25)...; levels = 1, colormap = :magma)
    test_ref("references/contourplot/gauss_nested.txt", @show_col(p))
end

@testset "function" begin
    p = @inferred contourplot(
        -3:0.01:3,
        -3:0.01:4,
        (x, y) -> (exp(-x^2 - y^2) + exp(-(x - 1)^2 - 2(y - 2)^2))^2,
    )
    test_ref("references/contourplot/function_contour.txt", @show_col(p))
end


text:
@recipe function f(a::SampleArray)
    xguide --> "s"
    yguide --> "amplitude"
    size --> (1000, 300)
#     ticks --> :native
    domain(a), data(a)
end


text:
# HamiltonianExpression is in fact an AutoMPO.
# But we implement some additional utils.

abstract type AbstractHamiltonianExpression end

export HamiltonianExpression

# ------------------------------------------------------
using ITensors:OpSum
struct HamiltonianExpression <: AbstractHamiltonianExpression
    ex :: OpSum
end

import Base.convert
convert(::Type{HamiltonianExpression},op::OpSum) = HamiltonianExpression(op)
HamiltonianExpression() = HamiltonianExpression(AutoMPO())

import Base.+
"""
    function +(a::HamiltonianExpression,b::HamiltonianExpression)
        Just add two hamiltonian expression together.
"""
function +(a::HamiltonianExpression,b::HamiltonianExpression)
    out = AutoMPO()
    out.data = append!(out.data,a.data,b.data)
    return out
end

"""
    function +(a::HamiltonianExpression,b)
        dispatch for AutoMPO native grammer support
"""
function +(a::HamiltonianExpression,b)
    return HamiltonianExpression(a.ex + b)
end


text:
module TestParsing

using Test
using DataFramesMeta
using Statistics

const ≅ = isequal

macro protect(x)
    esc(DataFramesMeta.get_column_expr(x))
end

@testset "Returning columnn identifiers" begin

    v_sym = @protect $[:a, :b]
    @test v_sym == [:a, :b]

    v_str = @protect $["a", "b"]
    @test v_str == ["a", "b"]

    qn = @protect :x
    @test qn == :x

    qn_protect = @protect $:x
    @test qn_protect == :x

    str_protect = @protect $"x"
    @test str_protect == "x"

    x = "a"
    sym_protect = @protect $x
    @test sym_protect === "a"

    v = ["a", "b"]
    v_protect = @protect $v
    @test v === v_protect

    b = @protect $(Between(:a, :b))
    @test b == Between(:a, :b)

    b = @protect $(begin 1 end)
    @test b == 1

    c = @protect cols(:a)
    @test c == :a

    i = @protect $1
    @test i == 1

    n = @protect "hello"
    @test n === nothing

    n = @protect 1
    @test n === nothing

    n = @protect begin x end
    @test n === nothing

    n = @protect x
    @test n === nothing
end

end # module


text:
module TestCycleGap

using Test
import GraphIdx

include("../src/cyclegap.jl")
include("demo3x7.jl")
n1, n2 = size(y)
lam = GraphIdx.Const(0.1)

@testset "Cycle Gap RotateDemo 3x7       " begin
    graph = GraphIdx.Grid.GridGraph(n1, n2)
    x = CycleGap.gaplas(copy(y), graph, lam, max_iter = 3, CycleGap.CycMem)
    @test size(x) == size(y)
end

end

max_stars_repo_name: EQt/treelas

max_issues_repo_name: EQt/treelas

max_forks_repo_name: EQt/treelas


text:
# Load Julia packages (libraries) needed for clip

using StatisticalRethinking

ProjDir = @__DIR__

# ### snippet 5.1

df = DataFrame!(CSV.File(rel_path("..", "data", "WaffleDivorce.csv"), delim=';'))
scale!(df, [:Marriage, :MedianAgeMarriage, :Divorce])

# Define the Stan language model

m_5_3 = "
data {
  int N;
  vector[N] divorce_s;
  vector[N] marriage_s;
  vector[N] medianagemarriage_s;
}
parameters {
  real a;
  real bA;
  real bM;
  real<lower=0> sigma;
}
model {
  vector[N] mu = a + + bA * medianagemarriage_s + bM * marriage_s;
  a ~ normal( 0 , 0.2 );
  bA ~ normal( 0 , 0.5 );
  bM ~ normal( 0 , 0.5 );
  sigma ~ exponential( 1 );
  divorce_s ~ normal( mu , sigma );
}
";

# Define the SampleModel
m5_3s = SampleModel("m5.3", m_5_3);

# Input data

m5_3_data = Dict(
  "N" => size(df, 1),
  "divorce_s" => df[:, :Divorce_s],
  "marriage_s" => df[:, :Marriage_s],
  "medianagemarriage_s" => df[:, :MedianAgeMarriage_s]
);

# Sample using cmdstan

rc = stan_sample(m5_3s, data=m5_3_data);

if success(rc)

  # Describe the draws

  dfa3 = read_samples(m5_3s; output_format=:dataframe)

  # Rethinking results

  rethinking_results = "
           mean   sd  5.5% 94.5%
    a      0.00 0.10 -0.16  0.16
    bM    -0.07 0.15 -0.31  0.18
    bA    -0.61 0.15 -0.85 -0.37
    sigma  0.79 0.08  0.66  0.91
  ";

  Particles(dfa3)

end


text:
using Gadfly
using DataFrames
using VML

# include("../utility.jl")

function main()

    N = [1,2,3,4,5,6,7,8,9,
         10,20,30,40,50,60,70,80,90,100,200,300,400,500,750,
         1_000,2_500,5_000,7_500,10_000,100_000,1_000_000]

    df = DataFrame(N=Int64[],ALLOC=Float64[],JULIA=Float64[],MKL=Float64[],MKL2=Float64[])

    @inbounds for i in 1:length(N)

        T = Float32
        n = N[i]
        a = rand(T,n)
        # println(n)

        Vector{T}(n)
        sqrt(a)
        VML.sqrt(a)
        VML.sqrt!(a,a)

        if n <= 1000
            gc()
            t1 = (@elapsed for k in 1:1000 Vector{T}(n) end)/1000
            t2 = (@elapsed for k in 1:1000 sqrt(a) end)/1000
            # t3 = (@elapsed for k in 1:1000 VML.sqrt(a) end)/1000
            # t4 = (@elapsed for k in 1:1000 VML.sqrt!(a,a) end)/1000
        end

        if n > 1000
            gc()
            t1 = mean([@elapsed Vector{T}(n) for k in 1:1000])
            t2 = mean([@elapsed sqrt(a) for k in 1:1000])
            # t3 = mean([@elapsed VML.sqrt(a) for k in 1:1000])
            # t4 = mean([@elapsed VML.sqrt!(a,a) for k in 1:1000])
        end

        push!(df,[n,t1*(1e9/n),t2*(1e9/n),0.0,0.0])
        # push!(df,[n,t1*(1e9/n),t2*(1e9/n),t3*(1e9/n),t4*(1e9/n)])
    end
    df
end

df = main()
println(df)
# @code_warntype main()

layers = Layer[]

df = DataFrame() ; df[:N], df[:CPU] = main1()
push!(layers,layer(df, x="N", y="CPU", Geom.line, Theme(default_color=color("blue")))[1])

df = DataFrame() ; df[:N], df[:CPU] = main2()
push!(layers,layer(df, x="N", y="CPU", Geom.line, Theme(default_color=color("yellow")))[1])

p = Gadfly.plot(layers,
                # Scale.y_log10,
                Scale.x_log10,
                Guide.xlabel("n-element vector"),
                Guide.ylabel("CPU time in nsec/element"),
                Guide.title("CPU time with n elements"),
                Guide.manual_color_key("",
                ["ALLOC", "Base.sqrt"],
                ["blue","yellow"])
                )

path = Pkg.dir("MKL") * "/benchmark/sqrt_cpu_regarding_n/"
draw(PNG(path*"sqrt_yichao.png", 20cm, 20cm), p)

df = main3();
println(df)

a = rand(2)
t = ibench2(sqrt,(a,),0.1)


text:
function exp!(X::IntervalBox)
    x, y = X

    y = y ∩ exp(x)
    x = x ∩ log(y)

    return IntervalBox(x, y)
end

"""
Reverse function for `exp`.
"""
function exp_rev(y::Interval, x::Interval)
    y_new = y ∩ (0..∞)
    x_new = x ∩ log(y)
    return y_new, x_new
end

"""
Reverse function for `exp2`.
"""
function exp2_rev(y::Interval, x::Interval)
    y_new = y ∩ (0..∞)
    x_new = x ∩ log2(y)

    return y_new, x_new
end

"""
Reverse function for `exp10`.
"""
function exp10_rev(y::Interval, x::Interval)
    y_new = y ∩ (0..∞)
    x_new = x ∩ log10(y)

    return y_new, x_new
end

"""
Reverse function for `expm1`.
"""
function expm1_rev(y::Interval, x::Interval)
    y_new = y ∩ (-1.0..∞)
    x_new = x ∩ log1p(y)

    return y_new, x_new
end

function log!(X::IntervalBox)  # y = log(x)
    x, y = X

    x = x ∩ exp(y)
    y = y ∩ log(x)

    return IntervalBox(x, y)
end

"""
Reverse function for `log`: ``y = \\log(x)``
"""
function log_rev(y::Interval, x::Interval)
    x_new = x ∩ exp(y)

    return y, x_new
end

"""
Reverse function for `log2`: ``y = \\log2(x)``
"""
function log2_rev(y::Interval, x::Interval)
    x_new = x ∩ exp2(y)

    return y, x_new
end

"""
Reverse function for `log10`: ``y = \\log10(x)``
"""
function log10_rev(y::Interval, x::Interval)
    x_new = x ∩ exp10(y)

    return y, x_new
end

"""
Reverse function for `log1p`: ``y = \\log1p(x)``
"""
function log1p_rev(y::Interval, x::Interval)
    x_new = x ∩ expm1(y)

    return y, x_new
end


text:
module Pinwheel
__precompile__(false)

export i, sq5, wheel, pinwheel

using Hecke
using Luxor
using StructEquality

using ...CoreDefs

Qx, x = QQ["x"]
K_ns, a = NumberField([x^2 - 5, x^2 + 1], ["√5", "i"])
K, nsmap = simple_extension(K_ns)
sq5 = inv(Pinwheel.nsmap)(a[1])
i = inv(Pinwheel.nsmap)(a[2])

ζ = (i - 2)//sq5
conj = complex_conjugation(K)

@def_structequal struct PinwheelElem <: DGroupElem
    rot_ζ :: Int
    rot_i :: Int
    refl :: Bool
    z :: nf_elem
end

# Hack because the nf_elem hash doesn't work for some reason
function Base.hash(x :: PinwheelElem, h::UInt)
    return hash((x.rot_ζ, x.rot_i, x.refl, string(x.z)), h)
end

function Base.:*(g :: PinwheelElem, h :: PinwheelElem)
    if !g.refl
        return PinwheelElem(
            g.rot_ζ + h.rot_ζ,
            mod(g.rot_i + h.rot_i, 4),
            h.refl,
            g.z + i^g.rot_i * (-ζ)^-g.rot_ζ * h.z
        )
    else
        return PinwheelElem(
            g.rot_ζ - h.rot_ζ,
            mod(g.rot_i - h.rot_i, 4),
            !h.refl,
            g.z + i^g.rot_i * (-ζ)^-g.rot_ζ * conj(h.z)
        )
    end
end
function CoreDefs.id(::Type{PinwheelElem})
    return PinwheelElem(0,0,0,K(0))
end
function Base.:*(g :: PinwheelElem, x :: nf_elem)
    if !g.refl
        return g.z + i^g.rot_i * (-ζ)^-g.rot_ζ * x
    else
        return g.z + i^g.rot_i * (-ζ)^-g.rot_ζ * conj(x)
    end
end
function CoreDefs.dilate(λ :: nf_elem, g :: PinwheelElem)
    return PinwheelElem(g.rot_ζ, g.rot_i, g.refl, λ*g.z)
end
function CoreDefs.id(::PinwheelElem)
    return PinwheelElem(0,0,0,K(0))
end
function Base.inv(p::PinwheelElem)
    return PinwheelElem(0,0,p.refl,K(0))*PinwheelElem(-p.rot_ζ,-p.rot_i,0,K(0))*PinwheelElem(0,0,0,-p.z)
end
function CoreDefs.embed_aff(g :: PinwheelElem)
    t = -atan(1/2)
    m = [0 -1; 1 0]^(g.rot_i)*[cos(t) sin(t); -sin(t) cos(t)]^g.rot_ζ

    z = convert(Complex{Float64}, evaluate(g.z, infinite_places(K)[2]))
    return [m[1,1], m[2,1], (-1)^g.refl*m[1,2], (-1)^g.refl*m[2,2], real(z), imag(z)]
end

struct PinwheelPTile end

function CoreDefs.draw(::PinwheelPTile, action)
    Luxor.poly([
        Point(0,0),
        Point(2,0),
        Point(0,1),
    ], close = true, action)
end

function wheel(rot_ζ, rot_i, refl, z)
    return PinwheelElem(rot_ζ, rot_i, refl, z) => PinwheelPTile()
end
function wheel()
    return wheel(0,0,0,K(0))
end

function color(p)
    return 1+Int(p[1].refl)
end

function pinwheel()
    sub  = Dict([
        (PinwheelPTile(), [
            wheel(-1,3,1,i*sq5-ζ)
            wheel(-1,0,1,sq5+2*ζ)
            wheel(-1,0,0,sq5+2*ζ)
            wheel(-1,0,1,2*sq5+2*ζ)
            wheel(-1,2,0,sq5*i-3*ζ)
        ])
    ])
    return SubSystem(sub, sq5)
end

const right = (1,0)
const ϑ = (0,1)

function vertices(t :: Pair{PinwheelElem, PinwheelPTile})
    return Dict([t[1]*K(0) => right, t[1]*K(1) => 2 .* right, t[1]*K(2) => ϑ, t[1]*(i) => right.-ϑ])
end

function precollar_in(tiling :: Dict, g :: PinwheelElem)
    gvertices = vertices(g => tiling[g])
    precollar =  []
    for t in tiling
        if !isempty(keys(vertices(t)) ∩ keys(gvertices))
            push!(precollar,t)
        end
    end
    return Dict(precollar)
end

function CoreDefs.is_interior(tiling :: Dict, g :: PinwheelElem)
    angle_sums = Dict(g => (0,0) for k in keys(vertices(g => tiling[g])))
    for tile in tiling
        for (vertex,angle) in vertices(tile)
            if vertex ∈ keys(angle_sums)
                angle_sums[vertex] = angle_sums[vertex] .+ angle
            end
        end
    end

    for (vertex, angle) in angle_sums
        if (angle[1] % 4 ≠ 0) || angle[2] ≠ 0
            return false
        end
    end
    return true
end

function CoreDefs.collar_in(tiling :: Dict, g :: PinwheelElem)
    precollar =  precollar_in(tiling , g)
    if !is_interior(precollar, g)
        throw(UnrecognizedCollar)
    else
        return precollar
    end
end

end


text:
using Win32API
using Base.Test

include("win32api_tests.jl")


text:
"""
    dataset(name::AbstractString)

Return a single dataset by its name.
"""
function dataset(name::AbstractString)
    q = datasets(Pair("name", name))
    return isequal(1)(length(q)) ? only(q) : nothing
end


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

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

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

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

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

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

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

    for nit=1:iter

        update_R!(X, R, Λ, T)

        Q₀ .= 0
        Q₁ .= 0

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

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

        if debug println(nit) end
    end

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


text:
using GeneratorsX
@generator function f(xs)
    for y in xs
        for x in y
            @yield x
        end
    end
end
collect(f([[1], [2, 3], [4, 5]]))


text:
STYLES = Dict(
    Initial => () -> TUI.Crayon(),
    Queued => () -> TUI.Crayon(),
    Running => () -> TUI.Crayon(foreground = :yellow),
    Passed => () -> TUI.Crayon(foreground = :green),
    Errored => () -> TUI.Crayon(foreground = :red),
    Failed => () -> TUI.Crayon(foreground = :red,),
)

function tuiloop(results; timeout = 3600, tsleep = 1/120)
    TUI.initialize()
    TUI.hide_cursor()
    terminal = TUI.Terminal()
    t = time()
    state = frame(terminal, results)
    while (time() - t) < timeout
        try
            state = frame(terminal, results, state)
        catch e
            if e isa InterruptException
                break
            else
                rethrow()
            end
        end
        sleep(tsleep)
    end
    TUI.cleanup()
    TUI.reset()
    TUI.disable_raw_mode()
end

function frame(terminal, results, state = (scroll = 1, selection = 1))
    scroll, selection = state
    w, h = TUI.terminal_size()

    # select list
    words = [TUI.Word(r.name, STYLES[r.state]()) for r in results]
    b1 = TUI.Block(title = "Packages")
    l_elem = TUI.SelectableList(
        b1,
        words,
        state.scroll,
        state.selection,
    )

    # test result box
    result = results[selection]
    #text = word.(split(results[selection].output))

    b2 = TUI.Block(title = "Test results")
    text_elem = Text(b2, resulttext(results[selection]))

    # layout and draw

    l_width = maximum(length.([r.name for r in results])) + 5
    l_rect = TUI.Rect(1, 1, l_width, h)
    TUI.draw(terminal, l_elem, l_rect)

    r2 = TUI.Rect(l_width + 3, 1, w - (l_width + 3), h - 1)
    TUI.draw(terminal, text_elem, r2)

    TUI.flush(terminal)

    # handle keyboard events
    c = TUI.get_event(terminal)
    if (c == 'q')
        throw(InterruptException())
    elseif (c == 'j')
        if selection < length(words)
            selection += 1
        end
    elseif (c == 'k')
        if selection > 1
            selection -= 1
        end
    end

    return (; scroll, selection)
end

function resulttext(result::TestResult)
    s = result.state
    if s === Running
        "Currently running..."
    elseif s === Queued
        "Queued and waiting to test"
    elseif s === Passed
        "TESTS PASSED\n\n" * result.output
    elseif s === Failed
        "TESTS FAILED\n\n" * result.output
    elseif s === Errored
        "TESTS ERRORED\n\n" * result.output
    elseif s === Initial
        "..."
    end
end


text:
@symbol_func function L_P(cur_reactor::AbstractReactor)
  cur_L = K_LP(cur_reactor)

  cur_L *= cur_reactor.R_0

  cur_L
end


text:
using Unums
using Base.Test

import Unums: Unum00, Ubound00, Ubound22, Ubound34

for U in (Ubound00, Ubound22, Ubound34, Unum00)
    @test U(0) + U(1) == U(1)
    @test U(0) - U(1) == U(-1)
    @test U(0) + U(2) == U(2)
    @test U(0) - U(2) == U(-2)
    @test U(1) + U(1) == U(2)
    @test U(1) - U(1) == U(0)
    @test U(1) + U(2) == U(3)
    @test U(1) - U(2) == U(-1)
    @test U(0) + U(0.5) == U(0.5)
    @test U(0) - U(0.5) == U(-0.5)
    @test U(1) + U(0.5) == U(1.5)
    @test U(1) - U(0.5) == U(0.5)

    @test U(1)*U(1) == U(1)
    @test U(1)*U(2) == U(2)
    @test U(1)/U(2) == U(0.5)
    @test U(1)/U(-2) == U(-0.5)
    @test sqrt(U(1)) == U(1)
    @test sqrt(U(2)) == U(sqrt(2))
    @test U(2)^4 == U(2^4)

    if U <: Unums.Ubound
        @test U(-1,1)^2 == U(0,1)
    end
end

for U in (Unum22, Ubound22)
    u = convert(U, 2.1)
    lo,hi,olo,ohi = limits(u)
    @test lo == 2
    @test hi == 2.125
    @test olo == true
    @test ohi == true
end

@test Unum22(2) in Unum22(2)
@test Unum22(2) in Ubound22(2,3)
@test !(Unum22(2) in Unum22(2.1))

@test isfinite(Unum22(2))
@test isfinite(Unum22(2.1))
@test !isfinite(Unum22(1e300))


text:
using Test

using MPI

# Call MPI.Init() before loading QCMaterial.mpi
MPI.Init()

using QCMaterial.mpi

@testset "mpi.distribute" begin
    size = 11
    comm_size = 3

    #rank = 0
    @test mpi.distribute(size, comm_size, 0) == (1, 4)

    #rank = 1
    @test mpi.distribute(size, comm_size, 1) == (5, 4)

    #rank = 2
    @test mpi.distribute(size, comm_size, 2) == (9, 3)
end

max_stars_repo_name: Sakurairihito/QCMaterialNew

max_issues_repo_name: sakurairihito/QCMaterialNew

max_forks_repo_name: sakurairihito/QCMaterialNew