text:
import CUDAnative

struct Dual{T}
    u::T
    du::T
end

import Base: promote_rule, convert
import Base.+, Base.-, Base.*, Base./, Base.^
import Base.>, Base.<, Base.>=, Base.<=, Base.==, Base.!=
import Base.isapprox

Dual(c::T) where T<:Number = Dual(c, zero(c))
convert(::Type{Dual{T}}, c::T) where T<:Number = Dual(c, zero(c))
promote_rule(::Type{Dual{T}}, ::Type{T}) where T<:Number = Dual{T}

(+)(x::Dual) = Dual(x.u, x.du)
(+)(x::Dual{T}, y::Dual{T}) where T = Dual(x.u+y.u, x.du+y.du)
(+)(x::Dual{T}, c::Number) where T = (+)(promote(x,T(c))...)
(+)(c::Number, y::Dual{T}) where T = (+)(promote(T(c),y)...)

(-)(x::Dual) = Dual(-x.u, -x.du)
(-)(x::Dual{T}, y::Dual{T}) where T = Dual(x.u-y.u, x.du-y.du)
(-)(x::Dual{T}, c::Number) where T = (-)(promote(x,T(c))...)
(-)(c::Number, y::Dual{T}) where T = (-)(promote(T(c),y)...)

(*)(x::Dual{T}, y::Dual{T}) where T = Dual(x.u*y.u, x.u*y.du+y.u*x.du)
(*)(x::Dual{T}, c::Number) where T = (*)(promote(x,T(c))...)
(*)(c::Number, y::Dual{T}) where T = (*)(promote(T(c),y)...)

(/)(x::Dual{T}, y::Dual{T}) where T = Dual(x.u/y.u, (x.du*y.u-x.u*y.du)/sqr(y.u))
(/)(x::Dual{T}, c::Number) where T = (/)(promote(x,T(c))...)
(/)(c::Number, y::Dual{T}) where T = (/)(promote(T(c),y)...)

CUDAnative.sin(x::Dual) = Dual(CUDAnative.sin(x.u), x.du*CUDAnative.cos(x.u))
CUDAnative.sin_fast(x::Dual) = Dual(CUDAnative.sin_fast(x.u), x.du*CUDAnative.cos_fast(x.u))
CUDAnative.sinpi(x::Dual{T}) where T = Dual(CUDAnative.sinpi(x.u), T(pi)*x.du*CUDAnative.cospi(x.u))

CUDAnative.cos(x::Dual) = Dual(CUDAnative.cos(x.u), -x.du*CUDAnative.sin(x.u))
CUDAnative.cos_fast(x::Dual) = Dual(CUDAnative.cos_fast(x.u), -x.du*CUDAnative.sin_fast(x.u))
CUDAnative.cospi(x::Dual{T}) where T = Dual(CUDAnative.cospi(x.u), -T(pi)*x.du*CUDAnative.sinpi(x.u))

function CUDAnative.tan(x::Dual)
    t = CUDAnative.tan(x.u)
    return Dual(t, x.du*(one(x.u)+sqr(t)))
end

function CUDAnative.tan_fast(x::Dual)
    t = CUDAnative.tan_fast(x.u)
    return Dual(t, x.du*(one(x.u)+sqr(t)))
end

CUDAnative.acos(x::Dual) = Dual(CUDAnative.acos(x.u), -x.du*CUDAnative.rsqrt(one(x.u)-sqr(x.u)))
CUDAnative.asin(x::Dual) = Dual(CUDAnative.asin(x.u), x.du*CUDAnative.rsqrt(one(x.u)-sqr(x.u)))
CUDAnative.atan(x::Dual) = Dual(CUDAnative.atan(x.u), x.du/(one(x.u)+sqr(x.u)))

CUDAnative.atan2(x::Dual{T}, y::Dual{T}) where T = Dual(CUDAnative.atan2(x.u,y.u), x.du/(one(x.u)+sqr(x.u/y.u)))
CUDAnative.atan2(x::Dual{T}, c::Number) where T = CUDAnative.atan2(promote(x,T(c))...)
CUDAnative.atan2(c::Number, x::Dual{T}) where T = CUDAnative.atan2(promote(T(c),x)...)

CUDAnative.cosh(x::Dual) = Dual(CUDAnative.cosh(x.u), x.du*CUDAnative.sinh(x.u))
CUDAnative.sinh(x::Dual) = Dual(CUDAnative.sinh(x.u), x.du*CUDAnative.cosh(x.u))

function CUDAnative.tanh(x::Dual)
    t = CUDAnative.tanh(x.u)
    return Dual(t, x.du*(one(x.u)-sqr(t)))
end

CUDAnative.acosh(x::Dual) = Dual(CUDAnative.acosh(x.u), x.du*CUDAnative.rsqrt(sqr(x.u)-one(x.u)))
CUDAnative.asinh(x::Dual) = Dual(CUDAnative.asinh(x.u), x.du*CUDAnative.rsqrt(sqr(x.u)+one(x.u)))
CUDAnative.atanh(x::Dual) = Dual(CUDAnative.atanh(x.u), x.du*inverse(one(x.u)-sqr(x.u)))

function CUDAnative.exp(x::Dual) # = Dual(CUDAnative.exp(x.u), x.du*CUDAnative.exp(x.u))
     t = CUDAnative.exp(x.u)
     return Dual(t, x.du*t)
 end

function CUDAnative.exp_fast(x::Dual)
    t = CUDAnative.exp_fast(x.u)
    return Dual(t, x.du*t)
end

function CUDAnative.exp2(x::Dual{T}) where T
    t = CUDAnative.exp2(x.u)
    return Dual(t, T(0.6931471805599453)*x.du*t)
end

function CUDAnative.exp10(x::Dual{T}) where T
    t = CUDAnative.exp10(x.u)
    return Dual(t, T(2.302585092994046)*x.du*t)
end

function CUDAnative.expm1(x::Dual)
    t = CUDAnative.expm1(x.u)
    return Dual(t, x.du*(one(x.u)+t))
end

CUDAnative.log(x::Dual) = Dual(CUDAnative.log(x.u), x.du/x.u)
CUDAnative.log_fast(x::Dual) = Dual(CUDAnative.log_fast(x.u), x.du/x.u)
CUDAnative.log10(x::Dual{T}) where T = Dual(CUDAnative.log10(x.u), T(0.43429448190325176)*x.du/x.u)
CUDAnative.log10_fast(x::Dual{T}) where T = Dual(CUDAnative.log10_fast(x.u), T(0.43429448190325176)*x.du/x.u)
CUDAnative.log2(x::Dual{T}) where T = Dual(CUDAnative.log2(x.u), T(1.4426950408889634)*x.du/x.u)
CUDAnative.log2_fast(x::Dual{T}) where T = Dual(CUDAnative.log2_fast(x.u), T(1.4426950408889634)*x.du/x.u)
CUDAnative.log1p(x::Dual) = Dual(CUDAnative.log1p(x.u), x.du/(one(x.u)+x.u))

CUDAnative.abs(x::Dual) = Dual(CUDAnative.abs(x.u), x.u > zero(x.u) ? x.du : -x.du)
CUDAnative.saturate(x::Dual) = Dual(CUDAnative.saturate(x.u), one(x.u) >= x.u >= zero(x.u) ? one(x.u) : zero(x.u))
CUDAnative.sqrt(x::Dual{T}) where T = Dual(CUDAnative.sqrt(x.u), T(0.5)*x.du/CUDAnative.sqrt(x.u))
CUDAnative.rsqrt(x::Dual{T}) where T = Dual(CUDAnative.rsqrt(x.u), T(-0.5)*x.du/(x.u*CUDAnative.sqrt(x.u)))
CUDAnative.cbrt(x::Dual{T}) where T = Dual(CUDAnative.cbrt(x.u), T(1/3)*x.du/sqr(CUDAnative.cbrt(x.u)))
CUDAnative.rcbrt(x::Dual{T}) where T = Dual(CUDAnative.rcbrt(x.u), T(-1/3)*x.du/sqr(sqr(CUDAnative.cbrt(x.u))))

function CUDAnative.hypot(x::Dual{T}, y::Dual{T}) where T
    t = CUDAnative.hypot(x.u, y.u)
    return Dual(t, (x.u*x.du+y.u*y.du)/t)
end

CUDAnative.hypot(x::Dual{T}, c::Number) where T = CUDAnative.hypot(promote(x,T(c))...)
CUDAnative.hypot(c::Number, x::Dual{T}) where T = CUDAnative.hypot(promote(T(c),x)...)

CUDAnative.erf(x::Dual{T}) where T = Dual(CUDAnative.erf(x.u), T(1.1283791670955126)*x.du*CUDAnative.exp(-sqr(x.u)))
CUDAnative.erfc(x::Dual{T}) where T = Dual(CUDAnative.erfc(x.u), T(-1.1283791670955126)*x.du*CUDAnative.exp(-sqr(x.u)))

function CUDAnative.erfinv(x::Dual{T}) where T
    t = CUDAnative.erfinv(x.u)
    Dual(t, T(0.8862269254527579)*x.du*CUDAnative.exp(sqr(t)))
end

function CUDAnative.erfcinv(x::Dual{T}) where T
    t = CUDAnative.erfcinv(x.u)
    Dual(t, T(-0.8862269254527579)*x.du*CUDAnative.exp(sqr(t)))
end

function CUDAnative.pow(x::Dual, p)
    t = CUDAnative.pow(x.u, p)
    return Dual(t, x.du*p*t/x.u)
end

function (^)(x::Dual, p)
    t = x.u ^ p
    return Dual(t, x.du*p*t/x.u)
end

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

CUDAnative.ceil(x::Dual) = Dual(CUDAnative.ceil(x.u), zero(x.u))
CUDAnative.floor(x::Dual) = Dual(CUDAnative.floor(x.u), zero(x.u))

CUDAnative.min(x::Dual{T}, y::Dual{T}) where T = CUDAnative.min(x.u, y.u)
CUDAnative.min(x::Dual{T}, c::Number) where T = CUDAnative.min(x.u, c)
CUDAnative.min(c::Number, x::Dual{T}) where T = CUDAnative.min(c, x.u)
CUDAnative.max(x::Dual{T}, y::Dual{T}) where T = CUDAnative.max(x.u, y.u)
CUDAnative.max(x::Dual{T}, c::Number) where T = CUDAnative.max(x.u, c)
CUDAnative.max(c::Number, x::Dual{T}) where T = CUDAnative.max(c, x.u)
CUDAnative.isfinite(x::Dual) = CUDAnative.isfinite(x.u)
CUDAnative.isinf(x::Dual) = CUDAnative.isinf(x.u)
CUDAnative.isnan(x::Dual) = CUDAnative.isnan(x.u)
CUDAnative.nearbyint(x::Dual) = CUDAnative.nearbyint(x.u)
CUDAnative.nextafter(x::Dual) = CUDAnative.nextafter(x.u)
CUDAnative.signbit(x::Dual) = CUDAnative.signbit(x.u)

(<)(x::Dual{T}, y::Dual{T}) where T = (x.u < y.u)
(<)(x::Dual{T}, c::Number) where T = (x.u < c)
(<)(c::Number, x::Dual{T}) where T = (c < x.u)

(>)(x::Dual{T}, y::Dual{T}) where T = (x.u > y.u)
(>)(x::Dual{T}, c::Number) where T = (x.u > c)
(>)(c::Number, x::Dual{T}) where T = (c > x.u)

(<=)(x::Dual{T}, y::Dual{T}) where T = (x.u <= y.u)
(<=)(x::Dual{T}, c::Number) where T = (x.u <= c)
(<=)(c::Number, x::Dual{T}) where T = (c <= x.u)

(>=)(x::Dual{T}, y::Dual{T}) where T = (x.u >= y.u)
(>=)(x::Dual{T}, c::Number) where T = (x.u >= c)
(>=)(c::Number, x::Dual{T}) where T = (c >= x.u)

(==)(x::Dual{T}, y::Dual{T}) where T = (x.u == y.u)
(==)(x::Dual{T}, c::Number) where T = (x.u == c)
(==)(c::Number, x::Dual{T}) where T = (c == x.u)

(!=)(x::Dual{T}, y::Dual{T}) where T = (x.u != y.u)
(!=)(x::Dual{T}, c::Number) where T = (x.u != c)
(!=)(c::Number, x::Dual{T}) where T = (c != x.u)

isapprox(x::Dual{T}, y::Dual{T}) where T = isapprox(x.u, y.u)
isapprox(x::Dual{T}, c::Number) where T = isapprox(x.u, c)
isapprox(c::Number, x::Dual{T}) where T = isapprox(c, x.u)

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


text:
using Yao, LinearAlgebra, YaoBlocks
import Random.shuffle

Ugate(nbit::Int, i::Int) = put(nbit, i=>chain(Rz(0), Rx(0), Rz(0)));
Ugate(i::Int) = put(i=>chain(Rz(0), Rx(0), Rz(0)));

struct InitStateMatrix
    m1::Matrix{ComplexF64}
    m2::Matrix{ComplexF64}
end

function build_state(mats::InitStateMatrix)
    # return zero_state(2) |> chain(2, put(1=>chain(Rz(1), Rx(2))), put(2=>chain(Rz(1), Rx(2))))
    return zero_state(3) |> chain(3, put(3, (2)=>matblock(mats.m1)), put(3, (3)=>matblock(mats.m2)))
end

# @show probs(build_state())

function overlap(mats::InitStateMatrix)
    return dot(zero_state(2) |> chain(2, put(2, (1)=>matblock(mats.m1)), put(2, (2)=>matblock(mats.m2)), control(1, 2=>X), put(1=>H)) |> probs,[1, 1, 1, -1])
end

function cost(trainX, trainY, circuit, c)
    actual_out = [dot(i |> build_state |> circuit |> probs, c) for i in trainX]
    return sum((trainY-actual_out).^2)/length(trainY)
end

function generate_gates_set(n::Int64)
    rtn = []
    for i in 1:n
        for j in 1:n
            if i != j
                push!(rtn, control(n, i, j=>X))
            end
        end
    end

    for i in 1:n
        push!(rtn, Ugate(n, i))
    end

    return rtn
end

# 检测生成的circuit是否合法
function validate_circuit(circuit)
    old = circuit[1]
    for i in 2:length(circuit)
        now = circuit[i]
        if isa(old, ControlBlock) && old == now
            return false
        end
        if (!isa(old, ControlBlock)) && (!isa(now, ControlBlock)) && occupied_locs(old) == occupied_locs(now)
            return false
        end
        old = now
    end

    if length(circuit) > 3 && length(param_indexes(circuit)) == 0
        return false
    end
    if occupied_locs(circuit[length(circuit)]) != (1, )
        return rand() < .2
    end

    return true
end

function param_indexes(circuit, indexes)
    rtn = []
    for i in indexes
        if nparameters(circuit[i]) > 0
            push!(rtn, i)
        end
    end
    return rtn
end

function param_indexes(circuit)
    return param_indexes(circuit, [ 1:length(circuit); ])
end

function generate_init_circuit(n::Int64, d::Int64)
    gate_set = generate_gates_set(n)
    rtn = chain(n)
    while length(rtn) < d
        push!(rtn, rand(gate_set))
        if !validate_circuit(rtn)
            pop!(rtn)
        end
    end
    params = rand(nparameters(rtn)) * 2π .- π
    dispatch!(rtn, params)

    return rtn
end

function random_modify!(circuit2, circuit_indexes, gate_set)
    for _ in 1:rand(1:5)
        index = rand(circuit_indexes)
        # @show index
        circuit2[index] = rand(gate_set)

        while !validate_circuit(circuit2)
            circuit2[index] = rand(gate_set)
        end
        params = rand(nparameters(circuit2[index])) * 2π .- π
        dispatch!(circuit2[index], params)
    end
end

function gradient(circuit, circuit_indexes, trainX, trainY, delta)
    n = nqubits(circuit)
    grad = zeros(Float64, nparameters(circuit))

    count = 1
    for i in 1:length(circuit)
        params = parameters(circuit[i])
        for k in 1:nparameters(circuit[i])
            if i in circuit_indexes
                params[k] += delta
                dispatch!(circuit[i], params)
                cost_pos = cost(trainX, trainY, circuit, c)

                params[k] -= 2 * delta
                dispatch!(circuit[i], params)
                cost_neg = cost(trainX, trainY, circuit, c)

                params[k] += delta
                dispatch!(circuit[i], params)

                # @show (cost_pos - cost_neg)

                grad[count] = (cost_pos - cost_neg)/(2 * delta)
            end
            count += 1
        end
    end
    return grad
end

using Flux.Optimise

c = [1, -1, 1, -1, 1, -1, 1, -1];
const cost_threshold = 2e-7
# circuit = build_circuit()
trainX = [InitStateMatrix(rand_unitary(2), rand_unitary(2)) for _ in 1:5]
testX = [InitStateMatrix(rand_unitary(2), rand_unitary(2)) for _ in 1:5]
trainY = map(overlap, trainX)
testY = map(overlap, testX)
opt = Descent(1.2)
d = 8
# opt = ADAM(0.03, (0.9, 0.999))

function train(circuit, circuit_indexes, trainX, trainY, testX, testY, cost_threshold)
    history = Float64[]
    last_cost = new_cost = 0
    for i in 1:60
        last_cost = new_cost
        new_cost = cost(trainX, trainY, circuit, c)
        # push!(history, new_cost)
        ps = parameters(circuit)
        # @show (i, ps, new_cost)
        if (abs(new_cost - last_cost) < new_cost * .005) || new_cost < cost_threshold
            # @show new_cost
            break
        end

        # Optimise only one gate each step
        # This prevents local-minimal
        for index in shuffle(circuit_indexes)
            Optimise.update!(opt, ps, gradient(circuit, [ index ], trainX, trainY, .001))
        end

        popdispatch!(circuit, ps)
        if i == 60
            printstyled(IOContext(stdout, :color => true), "Warning: failed to converge\n", color=:red)
        end
    end
    return (new_cost, history)
end

function compress_subsequence(annealing_param, circuit, trainX, cost_threshold)
    rtn = deepcopy(circuit)

    # generate random circuit_indexes
    index1 = rand([ 1:length(circuit) - 1; ])
    index2 = rand([ index1 + 1:min(index1 + 4, length(circuit)); ]) # try to compress 2-5 gates
    circuit_indexes = [ index1+1:index2; ]
    max_steps = 10 * (index2 - index1)

    gate_set = generate_gates_set(nqubits(circuit))

    rtn[index1] = Ugate(nqubits(rtn), rand([ 1:nqubits(rtn); ]))
    while !validate_circuit(rtn)
        rtn[index1] = Ugate(nqubits(rtn), rand([ 1:nqubits(rtn); ]))
    end

    newTrainY = [dot(i |> build_state |> circuit |> probs, c) for i in trainX]
    newTestY = [dot(i |> build_state |> circuit |> probs, c) for i in testX]
    train(rtn, param_indexes(rtn, circuit_indexes), trainX, newTrainY, testX, newTestY, cost_threshold)
    best_cost = cost(testX, newTestY, rtn, c)
    step = 0

    while best_cost > cost_threshold && step < max_steps
        step += 1;

        circuit_backup = deepcopy(rtn)
        random_modify!(rtn, circuit_indexes, gate_set)
        train(rtn, param_indexes(rtn, circuit_indexes), trainX, newTrainY, testX, newTestY, cost_threshold)
        new_cost = cost(testX, newTestY, rtn, c)

        accept_prob = exp(-annealing_param * (new_cost - best_cost))
        if new_cost < best_cost
            # @show (step, new_cost)
            best_cost = new_cost
            continue
        elseif rand() < accept_prob
            # accept
            # best_cost = new_cost
            # println("accepted")
        else
            # println("fallback")
            rtn = circuit_backup
        end
        # @show (step, best_cost, accept_prob)
    end

    sucs = best_cost < cost_threshold
    return (rtn, sucs)
end

function find_algo(annealing_param, d::Int64, trainX, trainY, max_steps::Int64, cost_threshold)
    gate_set = generate_gates_set(3)
    circuit = generate_init_circuit(3, d)
    train(circuit, param_indexes(circuit), trainX, trainY, testX, testY, cost_threshold)
    best_cost = cost(testX, testY, circuit, c)

    step = 0
    circuit_backup = circuit
    while best_cost > cost_threshold && step < max_steps
        step += 1;

        circuit_backup = deepcopy(circuit)
        random_modify!(circuit, [ 1:length(circuit); ], gate_set)

        train(circuit, param_indexes(circuit), trainX, trainY, testX, testY, cost_threshold)
        new_cost = cost(testX, testY, circuit, c)

        accept_prob = exp(-annealing_param * (new_cost - best_cost))
        if new_cost < best_cost
            @show (step, new_cost)
            best_cost = new_cost
            continue
        elseif rand() < accept_prob
            # accept
            # best_cost = new_cost
            println("accepted")
        else
            # println("fallback")
            circuit = circuit_backup
        end

        if mod(step, 5) == 0
            printstyled(IOContext(stdout, :color => true), "tring to compress...\n", color=:blue)

            (new_circuit, sucs) = compress_subsequence(annealing_param, circuit, trainX, cost_threshold)

            if sucs
                @show circuit
                printstyled(IOContext(stdout, :color => true), "compress succeeded!\n", color=:green)

                circuit = new_circuit
                best_cost = cost(testX, testY, circuit, c)

            else
                printstyled(IOContext(stdout, :color => true), "compress failed!\n", color=:cyan)
            end
            @show (step, d, best_cost)
            if mod(step, 100) == 0
                @show circuit
            end
        end
    end
    return (circuit, best_cost < cost_threshold)
end

function find_nearest_pi(x::Float64)
    rtn = mod(x, 2*π);
    rtn = rtn > π ? rtn - 2*π : rtn;
    rtn = π / 4 * round(4 * rtn / π)
end

function enhance_circuit_to_pi(circuit)
    params = parameters(circuit)
    params = map(find_nearest_pi, params)
    dispatch!(circuit, params)
end

# (tmp, history) = train(generate_init_circuit(3, 9), trainX, trainY, cost_threshold)
# using Plots
# function plot_history(history)
#     pyplot() # Choose a backend
#     fig1 = plot(history; legend=nothing)
#     title!("training history")
#     xlabel!("steps"); ylabel!("cost")
# end
# plot_history(history)

# annealing_param 必须是100000数量级
@show (circuit, ) = find_algo(70000.0, d, trainX, trainY, 500000, cost_threshold)
enhance_circuit_to_pi(circuit)
@show cost(trainX, trainY, circuit, c)

ext: jl


text:
# auxiliary functions:
# calculate plotting coordinates of nodes & edges for PhyloNetworks objects

"""
    getEdgeNodeCoordinates(net, useEdgeLength::Bool, useSimpleHybridLines::Bool)

Calculate coordinates for plotting later with Gadfly or RCall.

Actually modifies some (minor) attributes of the network,
as it calls `directEdges!` and `preorder!`.

output: tuple with the following elements, in which the order of
nodes corresponds to the order in `net.node`, and the order of
edges corresponds to that in `net.edge` (filtered to minor edges as needed).

1. `edge_xB`: x coordinate for the Beginning and ...
2. `edge_xE`: ...  End of each edge, in the same order as in `net.edge`
3. `edge_yB`: y coordinate for edges, Begin ...
4. `edge_yE`: ... and End.
   Each major edge is drawn as a single horizontal line. Minor hybrid edges are
   drawn as: a single diagonal segment if `useSimpleHybridLines` is true,
   or as 2 connected segments otherwise: one horizontal (whose length on the
   x axis can be used to represent the edge length), and the other diagonal to
   connect the horizontal segment to the child node.
   `edge_*` contains the coordinates for the horizontal segment only, which is
   reduced to a single point (Begin = End) when using "SimpleHybridLines".
   `minoredge_*` (see below) contains information for the diagonal segment.
   Agreed, `edge_yB` = `edge_yE` always (relic from before v0.3:
   no minoredge output back then, and simple diagonal lines only)
5. `node_x`: x and ...
6. `node_y`: ... y coordinate at the middle of the vertical bar that represents a node.
   The (or each) parent edge of the node connects to this middle point,
   but the node itself is drawn as a vertical bar connected to all it children edges.
   order: same as in `net.node`
7. `node_yB`: y coordinates of the Beginning and the ...
8. `node_yE`: ... End of the vertical bar representing each node.
   The x coordinate (Begin & End) of the end points of the vertical bar is
   the same as that of the mid-point, given by `node_x`.
9. `minoredge_xB`: x coordinate for the Beginning and ...
10. `minoredge_xE`: ... End of the diagonal segment of each minor hybrid edge,
   in the same order as in `filter(e -> !e.isMajor, net.edge)`.
11. `minoredge_yB`: y coordinate for the beginning and ...
12. `minoredge_yE`: ... end of the diagonal segment of each minor hybrid edge.
13-16. `xmin`, `xmax`, `ymin`, `ymax`: ranges for the x and y axes.
"""
function getEdgeNodeCoordinates(net::HybridNetwork, useEdgeLength::Bool, useSimpleHybridLines::Bool)
    try
        directEdges!(net)   # to update isChild1
    catch e
        if isa(e, PhyloNetworks.RootMismatch)
            e = PhyloNetworks.RootMismatch( e.msg * "\nPlease change the root, perhaps using rootatnode! or rootatedge!")
        end
        rethrow(e)
    end
    preorder!(net)       # to update net.nodes_changed: true pre-ordering

    # determine y for each node = y of its parent edge: post-order traversal
    # also [yB,yE] for each internal node: range of y's of all children nodes
    # y max is the numTaxa + number of minor edges
    ymin = 1.0;
    ymax = net.numTaxa
    if !useSimpleHybridLines
        ymax += sum(!e.isMajor for e in net.edge)
    end

    node_y  = zeros(Float64, net.numNodes) # order: in net.nodes, *!not in nodes_changed!*
    node_yB = zeros(Float64,net.numNodes) # min (B=begin) and max (E=end)
    node_yE = zeros(Float64,net.numNodes) #   of at children's nodes
    edge_yB = zeros(Float64,net.numEdges) # yE of edge = y of child node
    # set node_y of leaves: follow cladewise order
    # also sets edge_yB of minor hybrid edges
    nexty = ymax # first tips at the top, last at bottom
    cladewise_queue = copy(net.node[net.root].edge) # the child edges of root
    # print("queued the root's children's indices: "); @show queue
    while !isempty(cladewise_queue)
        cur_edge = pop!(cladewise_queue); # deliberate choice over shift! for cladewise order
        # increment spacing and add to node_y if leaf
        if getChild(cur_edge).leaf
            ni = findfirst(x->x===getChild(cur_edge), net.node)
            node_y[ni]  = nexty
            node_yB[ni] = nexty
            node_yE[ni] = nexty
            nexty -= 1
        end

        # only for new hybrid lines:
        # increment spacing and add to edge_yB if parent edge is minor
        if !cur_edge.isMajor && !useSimpleHybridLines
            edge_yB[findfirst(x->x===cur_edge, net.edge)] = nexty
            nexty -= 1
        end

        # push children edges if this is a major edge:
        if cur_edge.isMajor
            for e in getChild(cur_edge).edge
                if getParent(e) === getChild(cur_edge)# don't go backwards
                    push!(cladewise_queue, e)
                end
            end
        end
    end

    # set node_y of internal nodes: follow post-order
    for i=length(net.node):-1:1
        nn = net.nodes_changed[i]
        !nn.leaf || continue # previous loop took care of leaves
        ni = findfirst(x -> x===nn, net.node)
        node_yB[ni]=ymax; node_yE[ni]=ymin;
        minor_yB  = ymax; minor_yE  = ymin;
        nomajorchild=useSimpleHybridLines # only use this var if using simple hybrid lines
        for e in nn.edge
            if nn == PhyloNetworks.getParent(e) # if e = child of node
                if useSimpleHybridLines
                    # old simple hybrid lines
                    if e.isMajor || nomajorchild
                        cc = PhyloNetworks.getChild(e)
                        yy = node_y[findfirst(x -> x===cc, net.node)]
                        yy!==nothing || error("oops, child $(cc.number) has not been visited before node $(nn.number).")
                    end
                    if e.isMajor
                        nomajorchild = false # we found a child edge that is a major edge
                        node_yB[ni] = min(node_yB[ni], yy)
                        node_yE[ni] = max(node_yE[ni], yy)
                    elseif nomajorchild # e is minor edge, and no major found so far
                        minor_yB = min(minor_yB, yy)
                        minor_yE = max(minor_yE, yy)
                    end
                else
                    # new pretty hybrid lines
                    if e.isMajor
                        cc = PhyloNetworks.getChild(e)
                        child_y = node_y[findfirst(x -> x===cc, net.node)]
                        child_y!==nothing || error("oops, child $(cc.number) has not been visited before node $(nn.number).")
                    else
                        child_y = edge_yB[findfirst(x->x===e, net.edge)]
                    end
                    node_yB[ni] = min(node_yB[ni], child_y)
                    node_yE[ni] = max(node_yE[ni], child_y)
                end
            end
        end
        if nomajorchild # if all children edges are minor hybrid edges. If so: not level 1, not tree-child
            if minor_yB == minor_yE # one single child. jitter by 0.1 to make the plot readable
                minor_yB += (minor_yB < (ymax+ymin)/2 ? 0.1 : -0.1)
                minor_yE = minor_yB
            end
            node_yB[ni] = minor_yB
            node_yE[ni] = minor_yE
        end
        node_y[ni] = (node_yB[ni]+node_yE[ni])/2
        if nomajorchild #since the minor edges are leaving from the center of the node's y pos.
            node_yB[ni] = node_y[ni]
            node_yE[ni] = node_y[ni]
        end
    end

    # setting branch lengths for plotting
    elenCalculate = !useEdgeLength
    if (useEdgeLength)
        allBLmissing = true; nonBLmissing = true;
        for e in net.edge
            if (nonBLmissing && e.length==-1.0) nonBLmissing=false; end
            if (allBLmissing && e.length!=-1.0) allBLmissing=false; end
        end
        if (allBLmissing)
            println("All edge lengths are missing, won't be used for plotting.")
            elenCalculate = true
        end
        if (!nonBLmissing && !allBLmissing) # not all, but some are missing
            @warn "At least one non-missing edge length: plotting any missing length as 1.0"
        end
    end
    elen = Float64[] # edge lengths to be used for plotting. same order as net.edge.
    if (elenCalculate)
        # setting elen such that the age of each node = 1 + age of oldest child
        # (including minor hybrid edges): need true post-ordering.
        # calculating node ages first, elen will be calculated later.
        elen     = zeros(Float64,net.numEdges)
        node_age = zeros(Float64,net.numNodes)
        for i=length(net.node):-1:1 # post-order traversal
            if (net.nodes_changed[i].leaf) continue; end
            ni = findfirst(x -> x===net.nodes_changed[i], net.node)
            for e in net.nodes_changed[i].edge # loop over children only
                if net.nodes_changed[i] == (e.isChild1 ? e.node[2] : e.node[1])
                    node_age[ni] = max(node_age[ni], 1 +
                     node_age[findfirst(x -> x=== PhyloNetworks.getChild(e), net.node)])
                end
            end
        end
    else
        for e in net.edge
            push!(elen, (e.length==-1.0 ? 1.0 : e.length))
        end
    end

    # determine xB,xE for each edge: pre-order traversal, uses branch lengths
    # then x and yB,yE for each node: x=xE of parent edge
    xmin = 1.0; xmax=xmin
    node_x  = zeros(Float64,net.numNodes) # order: in net.nodes, *!not in nodes_changed!*
    edge_xB = zeros(Float64,net.numEdges) # min (B=begin) and max (E=end)
    edge_xE = zeros(Float64,net.numEdges) # xE-xB = edge length
    node_x[net.root] = xmin # root node: x=xmin=0
    for i=2:length(net.node)              # true pre-order, skipping the root (i=1)
        ni = findfirst(x -> x===net.nodes_changed[i], net.node)
        ei = nothing # index of major parent edge of current node
        for e in net.nodes_changed[i].edge
            if (e.isMajor && net.nodes_changed[i] == e.node[e.isChild1 ? 1 : 2]) # major parent edge
                ei = findfirst(x -> x===e, net.edge)
                break
            end
        end
        ei !== nothing || error("oops, could not find major parent edge of node number $ni.")
        edge_yB[ei] = node_y[ni]
        pni = findfirst(x -> x===PhyloNetworks.getParent(net.edge[ei]), net.node) # parent node index
        edge_xB[ei] = node_x[pni]
        if elenCalculate
            elen[ei] = node_age[pni] - node_age[ni]
        end
        edge_xE[ei] = edge_xB[ei] + elen[ei]
        node_x[ni] = edge_xE[ei]
    end
    edge_yE = copy(edge_yB) # true for tree and major edges

    # coordinates of the diagonal lines that connect hybrid edges with their targets
    minoredge_xB = Float64[]
    minoredge_xE = Float64[]
    minoredge_yB = Float64[]
    minoredge_yE = Float64[]

    for i=1:net.numEdges
        if (!net.edge[i].isMajor) # minor hybrid edges
            # indices of child and parent nodes
            cni = findfirst(x -> x===PhyloNetworks.getChild( net.edge[i]), net.node)
            pni = findfirst(x -> x===PhyloNetworks.getParent(net.edge[i]), net.node)

            edge_xB[i] = node_x[pni]
            edge_xE[i] = useSimpleHybridLines ? edge_xB[i] : (useEdgeLength ? edge_xB[i] + elen[i] : node_x[cni])

            if useSimpleHybridLines
                edge_yB[i] = node_y[pni]
            end
            edge_yE[i] = edge_yB[i]

            push!(minoredge_xB, edge_xE[i])
            push!(minoredge_yB, edge_yE[i])
            push!(minoredge_xE, node_x[cni])
            push!(minoredge_yE, node_y[cni])
            #@show i; @show net.edge[i]; @show pni; @show net.node[pni]; @show cni; @show net.node[cni]
        end
    end

    xmax = max(xmax, edge_xE...)

    #@show node_x;  @show node_yB; @show node_y;  @show node_yE
    #@show edge_xB; @show edge_xE; @show edge_yB; @show edge_yE
    return edge_xB, edge_xE, edge_yB, edge_yE,
           node_x, node_y, node_yB, node_yE,
           minoredge_xB, minoredge_xE, minoredge_yB, minoredge_yE,
           xmin, xmax, ymin, ymax
end

"""
    checkNodeDataFrame(net, nodeLabel)

Check data frame for node annotations:
- check that the data has at least 2 columns (if it has any)
- check that the first column has integers (to serve as node numbers)
- remove rows with no node numbers
- warning if some node numbers in the data are not in the network.
"""
function checkNodeDataFrame(net::HybridNetwork, nodeLabel::DataFrame)
    labelnodes = size(nodeLabel,1)>0
    if (labelnodes && (size(nodeLabel,2)<2 ||
            !(nonmissingtype(eltype(nodeLabel[!,1])) <: Integer)))
        @warn "nodeLabel should have 2+ columns, the first one giving the node numbers (Integer)"
        labelnodes = false
    end
    if labelnodes # remove rows with no node number, check if at least one row remains
        filter!(row->!ismissing(row[1]), nodeLabel)
        labelnodes = size(nodeLabel,1)>0
    end
    if labelnodes
      tmp = setdiff(nodeLabel[!,1], [n.number for n in net.node])
      if length(tmp)>0
        msg = "Some node numbers in the nodeLabel data frame are not found in the network:\n"
        for a in tmp msg *= string(" ",a); end
        @warn msg
      end
    end
    return(labelnodes, nodeLabel)
end

"""
    prepareNodeDataFrame(net, nodeLabel::DataFrame,
        showNodeNumber::Bool, showIntNodeLabel::Bool, labelnodes::Bool,
        node_x, node_y)

Make data frame for node annotation. `node_*` should be Float64 vectors.
`nodeLabel` should have columns as required by [`checkNodeDataFrame`](@ref)

Columns of output data frame:
- x, y: coordinates on the plots (from `node_*`)
- name: node name
- num: node number
- lab: node label
- lea: is leaf?
"""
function prepareNodeDataFrame(net::HybridNetwork, nodeLabel::DataFrame,
        showNodeNumber::Bool, showIntNodeLabel::Bool, labelnodes::Bool,
        node_x::Array{Float64,1}, node_y::Array{Float64,1})
    nrows = (showNodeNumber || showIntNodeLabel || labelnodes ? net.numNodes : net.numTaxa)
    ndf = DataFrame(:name => Vector{String}(undef,nrows),
        :num => Vector{String}(undef,nrows), :lab => Vector{String}(undef,nrows),
        :lea => Vector{Bool}(  undef,nrows), :x => Vector{Float64}( undef,nrows),
        :y => Vector{Float64}( undef,nrows), copycols=false)
    j=1
    for i=1:net.numNodes
    if net.node[i].leaf  || showNodeNumber || showIntNodeLabel || labelnodes
        ndf[j,:name] = net.node[i].name
        ndf[j,:num] = string(net.node[i].number)
        if labelnodes
          jn = findfirst(isequal(net.node[i].number), nodeLabel[!,1])
          ndf[j,:lab] = (jn===nothing || ismissing(nodeLabel[jn,2]) ? "" :  # node label not in table or missing
            (nonmissingtype(eltype(nodeLabel[!,2])) <: AbstractFloat ?
              @sprintf("%0.3g",nodeLabel[jn,2]) : string(nodeLabel[jn,2])))
        end
        ndf[j,:lea] = net.node[i].leaf # use this later to remove H? labels
        ndf[j,:y] = node_y[i]
        ndf[j,:x] = node_x[i]
        j += 1
    end
    end
    # @show ndf
    return(ndf)
end

"""
    prepareEdgeDataFrame(net, edgeLabel::DataFrame, mainTree::Bool,
        edge_xB, edge_xE, edge_yB, edge_yE,
        minoredge_xB, minoredge_xE, minoredge_yB, minoredge_yE)

Check data frame for edge annotation.
`edge_*`: Float64 vectors giving the coordinates for the beginning and end of edges.
Return data frame with columns
- x, y: coordinates on the plots
- len: node name
- gam: gamma (inheritance value)
- num: node number
- lab: node label
- hyb: is hybrid?
- min: is minor?
"""
function prepareEdgeDataFrame(net::HybridNetwork, edgeLabel::DataFrame, mainTree::Bool,
        edge_xB::Array{Float64,1}, edge_xE::Array{Float64,1},
        edge_yB::Array{Float64,1}, edge_yE::Array{Float64,1},
        minoredge_xB::Array{Float64,1}, minoredge_xE::Array{Float64,1},
        minoredge_yB::Array{Float64,1}, minoredge_yE::Array{Float64,1})
    nrows = net.numEdges - (mainTree ? length(minoredge_xB) : 0)
    edf = DataFrame(:len => Vector{String}(undef,nrows),
        :gam => Vector{String}(undef,nrows), :num => Vector{String}(undef,nrows),
        :lab => Vector{String}(undef,nrows), :hyb => Vector{Bool}(undef,nrows),
        :min => Vector{Bool}(  undef,nrows), :x => Vector{Float64}(undef,nrows),
        :y  => Vector{Float64}(undef,nrows), copycols=false)
    labeledges = size(edgeLabel,1)>0
    if (labeledges && (size(edgeLabel,2)<2 ||
            !(nonmissingtype(eltype(edgeLabel[!,1])) <: Integer)))
        @warn "edgeLabel should have 2+ columns, the first one giving the edge numbers (Integer)"
        labeledges = false
    end
    if labeledges # remove rows with no edge number and check if at least one remains
        filter!(row->!ismissing(row[1]), edgeLabel)
        labeledges = size(edgeLabel,1)>0
    end
    if labeledges
      tmp = setdiff(edgeLabel[!,1], [e.number for e in net.edge])
      if length(tmp)>0
        msg = "Some edge numbers in the edgeLabel data frame are not found in the network:\n"
        for a in tmp msg *= string(" ",a); end
        @warn msg
      end
    end
    j=1   # index of row in edf
    imh=1 # index of minor hybrid edge in filter(ee -> !ee.isMajor, net.edge)
    for i = 1:length(net.edge)
        ee = net.edge[i]
        # skip the edge if it's minor and we only want the main tree:
        mainTree && !ee.isMajor && continue
        edf[j,:len] = (ee.length==-1.0 ? "" : @sprintf("%0.3g",ee.length))
        # @sprintf("%c=%0.3g",'γ',ee.length)
        edf[j,:gam] = (ee.gamma==-1.0  ? "" : @sprintf("%0.3g",ee.gamma))
        edf[j,:num] = string(ee.number)
        if labeledges
            je = findfirst(isequal(ee.number), edgeLabel[!,1])
            edf[j,:lab] = (je===nothing || ismissing(edgeLabel[je,2]) ? "" :  # edge label not found in table
            (nonmissingtype(eltype(edgeLabel[!,2])) <: AbstractFloat ?
                @sprintf("%0.3g",edgeLabel[je,2]) : string(edgeLabel[je,2])))
        end
        edf[j,:hyb] = ee.hybrid
        edf[j,:min] = !ee.isMajor
        if ee.isMajor
            edf[j,:y] = (edge_yB[i] + edge_yE[i])/2
            edf[j,:x] = (edge_xB[i] + edge_xE[i])/2
        else
            edf[j,:y] = (minoredge_yB[imh] + minoredge_yE[imh])/2
            edf[j,:x] = (minoredge_xB[imh] + minoredge_xE[imh])/2
            imh += 1
        end
        j += 1
    end
    # @show edf
    return labeledges, edf
end

max_stars_repo_name: cecileane/PhyloPlots

max_issues_repo_name: cecileane/PhyloPlots

max_forks_repo_name: cecileane/PhyloPlots