text:
# This code is part of QuantumCircuits.
#
# (C) Copyright Rafał Pracht 2022.
#
# This code is licensed under the Apache License, Version 2.0. You may
# obtain a copy of this license in the LICENSE.txt file in the root directory
# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
#
# Any modifications or derivative works of this code must retain this
# copyright notice, and modified files need to carry a notice indicating
# that they have been altered from the originals.

module Circuit

using QuantumCircuits.QCircuits.QBase
using QuantumCircuits.QCircuits.Gates
using QuantumCircuits.QCircuits.Gates: ParamT, appendparams!, appendRandParams!
using QuantumCircuits.QCircuits.ComplexGates
using QuantumCircuits.QCircuits.ComplexGates: U4_params
using QuantumCircuits.QCircuits.OtherGates
using QuantumCircuits.QCircuits.Instructions
using QuantumCircuits.QCircuits.Registers
using QuantumCircuits.QCircuits.Qiskit
using QuantumCircuits.QCircuits.Graph
using QuantumCircuits.QCircuits.Math
using QuantumCircuits.QCircuits.Qiskit: IS_QISKIT, NoQiskitError

using LinearAlgebra: I

import QuantumCircuits.QCircuits.QBase: add!, tomatrix, setparameters!, simplify,
       standardGateError, decompose, measure!, bindparameters!
import Base.show

export QCircuit, getparameters, getRandParameters, toString, setClassicalRegister!

"Nothing function"
const nop = () -> nothing

"Quantum circuit."
mutable struct QCircuit <: QuantumCircuit
    qubits::Integer
    qRegisters::Vector{QuantumRegister}
    cRegisters::Vector{ClassicalRegister}
    dcg::DirectedGraph
    has_code::Bool
    code::Vector{QuantumGate}
    vqubits::Vector{Qubit}
    vcbits::Vector{Cbit}
    measures::Vector{Pair{Qubit, Cbit}}
    measures_matrix::Matrix{Float64}

    # gates
    x::Function
    sx::Function
    y::Function
    z::Function
    h::Function
    cx::Function
    s::Function
    sdg::Function
    t::Function
    tdg::Function
    u::Function
    u3::Function
    rx::Function
    ry::Function
    rz::Function
    rzx::Function
    u4::Function
    barrier::Function
    measure::Function

    function QCircuit(qRegs::Vector{QuantumRegister}, cRegs::Vector{ClassicalRegister})
        n = sum([length(i) for i in qRegs])

        qc = new(n, qRegs, cRegs, DirectedGraph(n), false, QuantumGate[], Qubit[], ClassicalRegister[], Pair{Qubit, Cbit}[], eye(2^n), nop)

        # assign the qregister to circuit
        for r in qRegs
            append!(qc.vqubits, r.bits)
        end
        for (i, q) in enumerate(qc.vqubits)
            setid(q, i - 1)
        end

        # and classical one
        for r in cRegs
            append!(qc.vcbits, r.bits)
        end
        for (i, q) in enumerate(qc.vcbits)
            setid(q, i - 1)
        end

        # add gates
        qc.x = q -> add!(qc, X, q)
        qc.sx = q -> add!(qc, Sx, q)
        qc.y = q -> add!(qc, Y, q)
        qc.z = q -> add!(qc, Z, q)
        qc.h = q -> add!(qc, H, q)
        qc.cx = (c, t) -> add!(qc, CX, c, t)
        qc.s = q -> add!(qc, S, q)
        qc.sdg = q -> add!(qc, Sd, q)
        qc.t = q -> add!(qc, T, q)
        qc.tdg = q -> add!(qc, Td, q)
        qc.u = (q, θ, ϕ, λ) -> add!(qc, U, q, θ, ϕ, λ)
        qc.u3 = (q, θ=ParameterT(rand()*2π), ϕ=ParameterT(rand()*2π), λ=ParameterT(rand()*2π)) -> add!(qc, U3, q, θ, ϕ, λ)
        qc.rx = (q, θ=ParameterT(rand()*2π)) -> add!(qc, Rx, q, θ)
        qc.ry = (q, θ=ParameterT(rand()*2π)) -> add!(qc, Ry, q, θ)
        qc.rz = (q, θ=ParameterT(rand()*2π)) -> add!(qc, Rz, q, θ)
        qc.rzx = (q1, q2, θ=ParameterT(rand()*2π)) -> add!(qc, Rzx, q1, q2, θ)
        qc.u4 = (q1, q2, params=[ParameterT(rand() * 2π) for i in 1:U4_params]) -> add!(qc, U4, q1, q2, params)
        qc.barrier = () -> add!(qc, Barrier(qc.vqubits))
        qc.measure = (q, c) -> measure!(qc, q, c)

        return qc
    end
end
QCircuit(n::Integer) = QCircuit([QuantumRegister(n)], [ClassicalRegister(n)])
QCircuit(qReg::QuantumRegister, cReg::ClassicalRegister) = QCircuit([qReg], [cReg])
QCircuit(reg::QuantumRegister) = QCircuit([reg], ClassicalRegister[])
QCircuit(regs::Vector{QuantumRegister}) = QCircuit(regs, [ClassicalRegister(sum([length(i) for i in regs]))])
QCircuit(regs::Vector{QuantumRegister}, cReg::ClassicalRegister) = QCircuit(regs, [cReg])
function QCircuit(qc::QCircuit)
    qregs = [QuantumRegister(length(r), r.name) for r in qc.qRegisters]
    cregs = [ClassicalRegister(length(r), r.name) for r in qc.cRegisters]

    QCircuit(qregs, cregs)
end

function getCode(c::QCircuit)
    if c.has_code
        return c.code
    else
        c.code = to_vector(c.dcg)
        c.has_code = true

        return c.code
    end
end

Base.:(==)(c1::QCircuit, c2::QCircuit) = c1.qubits == c2.qubits && getCode(c1) == getCode(c2)
Base.hash(c::QCircuit, h::UInt) = hash((c.qubits, getCode(code)), h)

"Add gate to circuit"
function add!(qc::QCircuit, gate::QuantumGate)
    add!(qc.dcg, gate)
    qc.has_code = false
    nothing
end
function add!(qc::QCircuit, gate::Type{T}, qubit::Qubit) where T <: QuantumGate
    @assert qubit in qc.vqubits "Qubits is out of circuit qubits range."

    add!(qc.dcg, gate(qubit))
    qc.has_code = false
    nothing
end
function add!(qc::QCircuit, gate::Type{T}, args::Tuple) where T <: QuantumGate
    add!(qc.dcg, gate(args...))
    qc.has_code = false
    nothing
end
function add!(qc::QCircuit, gate::Type{T}, qubit::Integer) where T <: QuantumGate
    @assert qubit < qc.qubits "Qubits $qubit is out of circuit qubits range $(qc.qubits)."

    add!(qc, gate, qc.vqubits[qubit + 1])

    nothing
end
function add!(qc::QCircuit, gate::Type{T}, qubits::Vector) where T <: QuantumGate
    for qubit in qubits
        add!(qc, gate, qubit)
    end
end
function add!(qc::QCircuit, gate::Type{T}, qubits::QuantumRegister) where T <: QuantumGate
    for qubit in qubits
        add!(qc, gate, qubit)
    end
end
function add!(qc::QCircuit, gate::Type{T}, qubit1::Qubit, qubit2::Qubit) where T <: QuantumGate
    @assert qubit1 in qc.vqubits "Qubits is out of circuit qubits range."
    @assert qubit2 in qc.vqubits "Qubits is out of circuit qubits range."

    add!(qc.dcg, gate(qubit1, qubit2))
    qc.has_code = false
    nothing
end
function add!(qc::QCircuit, gate::Type{T}, qubit1::Integer, qubit2::Integer) where T <: QuantumGate
    @assert qubit1 < qc.qubits "Qubits $qubit1 is out of circuit qubits range $(qc.qubits)."
    @assert qubit2 < qc.qubits "Qubits $qubit2 is out of circuit qubits range $(qc.qubits)."

    add!(qc, gate, qc.vqubits[qubit1 + 1], qc.vqubits[qubit2 + 1])

    nothing
end
function add!(qc::QCircuit, gate::Type{T}, qubit1::Qubit, qubit2::Qubit, params) where T <: QuantumGate
    @assert qubit1 in qc.vqubits "Qubits is out of circuit qubits range."
    @assert qubit2 in qc.vqubits "Qubits is out of circuit qubits range."

    if params == nothing
        add!(qc.dcg, gate(qubit1, qubit2))
    else
        add!(qc.dcg, gate(qubit1, qubit2, params))
    end
    qc.has_code = false
    nothing
end
function add!(qc::QCircuit, gate::Type{T}, qubit1::Integer, qubit2::Integer, params) where T <: QuantumGate
    @assert qubit1 < qc.qubits "Qubits $qubit1 is out of circuit qubits range $(qc.qubits)."
    @assert qubit2 < qc.qubits "Qubits $qubit2 is out of circuit qubits range $(qc.qubits)."

    add!(qc, gate, qc.vqubits[qubit1 + 1], qc.vqubits[qubit2 + 1], params)

    nothing
end

function add!(qc::QCircuit, gate::Type{T}, qubit::Qubit, θ) where T <: QuantumGate
    @assert qubit in qc.vqubits "Qubits is out of circuit qubits range."

    add!(qc.dcg, gate(qubit, θ))
    qc.has_code = false
    nothing
end
function add!(qc::QCircuit, gate::Type{T}, qubit::Integer, θ) where T <: QuantumGate
    @assert qubit < qc.qubits "Qubits $qubit is out of circuit qubits range $(qc.qubits)."

    add!(qc, gate, qc.vqubits[qubit + 1], θ)

    nothing
end
function add!(qc::QCircuit, gate::Type{T}, qubits::Vector, θ) where T <: QuantumGate
    for qubit in qubits
        add!(qc, gate, qubit, θ)
    end
end

function add!(qc::QCircuit, gate::Type{T}, qubit::Qubit, θ, ϕ, λ) where T <: QuantumGate
    @assert qubit in qc.vqubits "Qubits is out of circuit qubits range."

    add!(qc.dcg, gate(qubit, θ, ϕ, λ))
    qc.has_code = false
    nothing
end
function add!(qc::QCircuit, gate::Type{T}, qubit::Integer, θ, ϕ, λ) where T <: QuantumGate
    @assert qubit < qc.qubits "Qubits $qubit is out of circuit qubits range $(qc.qubits)."

    add!(qc, gate, qc.vqubits[qubit + 1], θ, ϕ, λ)

    nothing
end
function add!(qc::QCircuit, gate::Type{T}, qubits::QuantumRegister, θ, ϕ, λ) where T <: QuantumGate
    for qubit in qubits
        add!(qc, gate, qubit, ParameterT(θ), ParameterT(ϕ), ParameterT(λ))
    end
end
function add!(qc::QCircuit, gate::Type{T}, qubits::Vector, θ, ϕ, λ) where T <: QuantumGate
    for qubit in qubits
        add!(qc, gate, qubit, ParameterT(θ), ParameterT(ϕ), ParameterT(λ))
    end
end

function add!(qc::QCircuit, gates::Vector{T}) where T <: QuantumGate
    for g in gates
        add!(qc, g)
    end
    nothing
end

"Add measuers to circuit"
function measure!(qc::QCircuit, qubit::Qubit, cbit::Cbit, setMatrix::Bool=true)
    push!(qc.measures, qubit => cbit)

    if setMatrix
        setMeasureMatrix!(qc)
    end
end
"Add measuers to circuit"
function measure!(qc::QCircuit, qubit::Integer, cbit::Integer, setMatrix::Bool=true)
    @assert qubit < qc.qubits "Qubits $qubit is out of circuit qubits range $(qc.qubits)."
    @assert cbit < length(qc.vcbits) "Qubits $qubit is out of circuit qubits range $(length(qc.vcbits))."

    measure!(qc, qc.vqubits[qubit + 1], qc.vcbits[cbit + 1], setMatrix)
end
function measure!(qc::QCircuit, qubits::Vector{T}, cbits::Vector{T}) where T <: Integer
    n = length(qubits)
    @assert n == length(cbits) "The length of qubits and classical bits should be equal."

    for i in 1:n
        measure!(qc, qubits[i], cbits[i], false)
    end

    setMeasureMatrix!(qc)
end

function setMeasureMatrix!(qc::QCircuit)
    measuered = [(i-1, v) for (i, v) in enumerate(sort([getid(q)  for (q, c) in qc.measures]))]
    measuered_qubits = length(measuered)
    mes_matrix = zeros(2^measuered_qubits, 2^qc.qubits)
    for i in 0:(2^qc.qubits-1)
        bin = digits(i, base=2, pad=qc.qubits)

        # calculate new index
        ni = 0
        for (idx, m) in measuered
            if bin[m + 1] == 1
                ni += 2^idx
            end
        end
        mes_matrix[ni+1, i+1] = 1.0
    end

    qc.measures_matrix = mes_matrix
end

"Add the classical register to circuit."
function setClassicalRegister!(qc::QCircuit, cr::ClassicalRegister)
    @assert qc.cRegisters == ClassicalRegister[] "Unable to set register."
    @assert qc.vcbits == Cbit[] "Internal error :P"

    push!(qc.cRegisters, cr)
    append!(qc.vcbits, cr.bits)
    for (i, q) in enumerate(qc.vcbits)
        setid(q, i - 1)
    end
end

"Function to convert QCircuit to Qiskit circuit"
function toQiskit(circuit::QCircuit)
    if !IS_QISKIT
        throw(NoQiskitError("Unable to convert to Qiskit, qiskit is not installed."))
    end

    if isempty(circuit.cRegisters)
        qc = QiskitCircuit(circuit.qRegisters, nothing)
    else
        qc = QiskitCircuit(circuit.qRegisters, circuit.cRegisters)
    end

    for gate in getCode(circuit)
        add!(qc, gate)
    end

    for (q, c) in circuit.measures
        measure!(qc, q, c)
    end

    return qc
end

"Show method"
function Base.show(io::IO, qc::QCircuit)
    if IS_QISKIT
        print(io, toQiskit(qc))
    else
        print(io, "QCircuit")
    end
end

"Convert circuit to string"
function toString(qc::QCircuit)
    io = IOBuffer()
    show(io, qc)
    return String(take!(io))
end

"Function create new parameters vector"
function getparameters(qc::QCircuit)
    params = ParamT[]
    for gate in getCode(qc)
        appendparams!(params, gate)
    end

    return params
end

"Function create new random parameters vector"
function getRandParameters(qc::QCircuit)
    params = ParamT[]
    for gate in getCode(qc)
        appendRandParams!(params, gate)
    end

    return params
end

"Save the parameters to the quantum circuits."
function setparameters!(qc::QCircuit, params)
    i = 1
    for gate in getCode(qc)
        l = length(gate)
        if l > 0
            @views setparameters!(gate, params[i:(i+l-1)])
            i = i + l
        end
    end

    @assert i - 1 == length(params) "There is different number of parameters in circuit $(i-1) and in vector $(length(params))."
end

"Function bind the parameters, so now there only values."
function bindparameters!(qc::QCircuit)
    for gate in getCode(qc)
        bindparameters!(gate)
    end
end

function tomatrix(qc::QCircuit, params=nothing)
    i = 1
    unitary = I * (1.0 + 0im)
    for gate in getCode(qc)
        l = length(gate)
        if params != nothing && l > 0
            @views gate_unitary = tomatrix(qc.qubits, gate, params[i:(i+l-1)])
            # TODO gate_unitary = tomatrix(qc.qubits, gate, params[i:(i+l-1)])
            i = i + l
        else
            gate_unitary = tomatrix(qc.qubits, gate)
        end

        #println("GU: $gate_unitary")
        unitary = gate_unitary * unitary
        #println("U: $unitary")
    end
    @assert i - 1 == length(params) "There is different number of parameters in circuit $(i-1) and in vector $(length(params))."

    return unitary
end

function standardGateError(qc::QCircuit, params=nothing)
    err = 0.0
    i = 1
    for gate in getCode(qc)
        l = length(gate)
        if params != nothing && l > 0
            @views gerr = standardGateError(gate, params[i:(i+l-1)])
            i = i + l
        else
            gerr = standardGateError(gate)
        end

        err +=  gerr
    end

    #return err / (100 * length(qc.code))
    return err
end

function Base.inv(qc::QCircuit)
    newqc = QCircuit(qc)

    for gate in reverse(getCode(qc))
        newgate = inv(gate)
        add!(newqc, typeof(newgate), getArgs(newgate)...)
    end

    return newqc
end

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

function decompose(qc::QCircuit)
    newqc = QCircuit(qc)

    for gate in getCode(qc)
        newgates = decompose(gate)
        for newgate in newgates
            add!(newqc, typeof(newgate), getArgs(newgate)...)
        end
    end

    return newqc
end

function simplify(qc::QCircuit)
    newqc = QCircuit(qc)

    for gate in getCode(qc)
        newgates = simplify(gate)
        for newgate in newgates
            add!(newqc, typeof(newgate), getArgs(newgate)...)
        end
    end

    return newqc
end

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

function Base.append!(qc::QCircuit, apqc::QCircuit)
    @assert length(qc.measures) == 0 "Unable append to measured circuit."
    @assert length(apqc.measures) == 0 "Unable append measured circuit."
    @assert length(qc.vqubits) == length(apqc.vqubits) "Unable append, the qubits are different"

    for (q, aq) in zip(qc.vqubits, apqc.vqubits)
        @assert q == aq "There are difference in qubits."
    end

    for c in getCode(apqc)
        add!(qc, typeof(c), getArgs(c)...)
    end
end

end  # module Circuit


text:
########## Script ##########

function main(input_name)

    # open file and store content into an input string
    input = open(join(["inputs//", input_name]), "r") do io
        read(io, String)
    end

    # split input into a data vector
    data = String.(split(input, "\r\n"))

    # tables from day10 instructions
    pairs = Dict('('=>')', '['=>']', '{'=>'}', '<'=>'>')
    score1 = Dict(')'=>3, ']'=>57, '}'=>1197, '>'=>25137)
    score2 = Dict(')'=>1, ']'=>2, '}'=>3, '>'=>4)

    # part 1 and 2
    total_score1 = 0
    total_scores2 = Int[]
    for line in data

        opened_chars = Char[]
        corrupted = false
        for char in line
            if char in keys(pairs)
                append!(opened_chars, [char])
            elseif char in values(pairs)
                if char == pairs[opened_chars[end]]
                    pop!(opened_chars)
                else
                    total_score1 += score1[char]
                    corrupted = true
                    break
                end
            end
        end

        if !corrupted
            total_score2 = 0
            for char in reverse(opened_chars)
                total_score2 *= 5
                total_score2 += score2[pairs[char]]
            end
            append!(total_scores2, [total_score2])
        end

    end
    sort!(total_scores2)

    println("Part 1: ", total_score1)
    println("Part 2: ", total_scores2[(end + 1) ÷ 2])

end

main("day10.txt")

ext: jl


text:
# We will need to detect new function bodies, compare function bodies
# to see if they've changed, etc.  This has to be done "blind" to the
# line numbers at which the functions are defined.
#
# Now, we could just discard line numbers from expressions, but that
# would have a very negative effect on the quality of backtraces. So
# we keep them, but introduce machinery to compare expressions without
# concern for line numbers.

"""
A `RelocatableExpr` wraps an `Expr` to ensure that comparisons
between `RelocatableExpr`s ignore line numbering information.
This allows one to detect that two expressions are the same no matter
where they appear in a file.
"""
struct RelocatableExpr
    ex::Expr
end

const ExLike = Union{Expr,RelocatableExpr}

Base.convert(::Type{Expr}, rex::RelocatableExpr) = rex.ex
# Expr(rex::RelocatableExpr) = rex.ex   # too costly (inference invalidation)

Base.copy(rex::RelocatableExpr) = RelocatableExpr(copy(rex.ex))

# Implement the required comparison functions. `hash` is needed for Dicts.
function Base.:(==)(ra::RelocatableExpr, rb::RelocatableExpr)
    a, b = ra.ex, rb.ex
    if a.head == b.head
    elseif a.head == :block
        a = unwrap(a)
    elseif b.head == :block
        b = unwrap(b)
    end
    return a.head == b.head && isequal(LineSkippingIterator(a.args), LineSkippingIterator(b.args))
end

const hashrex_seed = UInt == UInt64 ? 0x7c4568b6e99c82d9 : 0xb9c82fd8
Base.hash(x::RelocatableExpr, h::UInt) = hash(LineSkippingIterator(x.ex.args),
                                              hash(x.ex.head, h + hashrex_seed))

function Base.show(io::IO, rex::RelocatableExpr)
    show(io, striplines!(copy(rex.ex)))
end

function striplines!(ex::Expr)
    if ex.head == :macrocall
        # for macros, the show method in Base assumes the line number is there,
        # so don't strip it
        args3 = [a isa ExLike ? striplines!(a) : a for a in ex.args[3:end]]
        return Expr(ex.head, ex.args[1], nothing, args3...)
    end
    args = [a isa ExLike ? striplines!(a) : a for a in ex.args]
    fargs = collect(LineSkippingIterator(args))
    return Expr(ex.head, fargs...)
end
striplines!(rex::RelocatableExpr) = RelocatableExpr(striplines!(rex.ex))

# We could just collect all the non-line statements to a Vector, but
# doing things in-place will be more efficient.

struct LineSkippingIterator
    args::Vector{Any}
end

Base.IteratorSize(::Type{LineSkippingIterator}) = Base.SizeUnknown()

function Base.iterate(iter::LineSkippingIterator, i=0)
    i = skip_to_nonline(iter.args, i+1)
    i > length(iter.args) && return nothing
    return (iter.args[i], i)
end

function skip_to_nonline(args, i)
    while true
        i > length(args) && return i
        ex = args[i]
        if isa(ex, Expr) && ex.head == :line
            i += 1
        elseif isa(ex, LineNumberNode)
            i += 1
        elseif isa(ex, Pair) && (ex::Pair).first == :linenumber     # used in the doc system
            i += 1
        elseif isa(ex, Base.RefValue) && !isdefined(ex, :x)         # also in the doc system
            i += 1
        else
            return i
        end
    end
end

function Base.isequal(itera::LineSkippingIterator, iterb::LineSkippingIterator)
    # We could use `zip` here except that we want to insist that the
    # iterators also have the same length.
    reta, retb = iterate(itera), iterate(iterb)
    while true
        reta === nothing && retb === nothing && return true
        (reta === nothing || retb === nothing) && return false
        vala, ia = reta
        valb, ib = retb
        if isa(vala, Expr) && isa(valb, Expr)
            vala.head == valb.head || return false
            isequal(LineSkippingIterator(vala.args), LineSkippingIterator(valb.args)) || return false
        elseif isa(vala, Symbol) && isa(valb, Symbol)
            # two gensymed symbols do not need to match
            sa, sb = String(vala), String(valb)
            (startswith(sa, '#') && startswith(sb, '#')) || isequal(vala, valb) || return false
        elseif isa(vala, Number) && isa(valb, Number)
            vala === valb || return false    # issue #233
        else
            isequal(vala, valb) || return false
        end
        reta, retb = iterate(itera, ia), iterate(iterb, ib)
    end
end

const hashlsi_seed = UInt == UInt64 ? 0x533cb920dedccdae : 0x2667c89b
function Base.hash(iter::LineSkippingIterator, h::UInt)
    h += hashlsi_seed
    for x in iter
        if x isa Expr
            h += hash(LineSkippingIterator(x.args), hash(x.head, h + hashrex_seed))
        elseif x isa Symbol
            xs = String(x)
            if startswith(xs, '#')  # all gensymmed symbols are treated as identical
                h += hash("gensym", h)
            else
                h += hash(x, h)
            end
        elseif x isa Number
            h += hash(typeof(x), hash(x, h))
        else
            h += hash(x, h)
        end
    end
    h
end


text:
using Test
using TestSetExtensions
using LinearAlgebra
using Random
using Qaintessent

##==----------------------------------------------------------------------------------------------------------------------

@testset ExtendedTestSet "statevector" begin

    @testset "test statevector int constructor" begin
        N = rand(1:10)
        s = Statevector(N)
        @test s.N == N
        @test s.state == zeros(ComplexF64, 2^N)
    end

    @testset "test statevector vector constructor" begin
        N = rand(1:10)
        ψ = randn(ComplexF64, 2^N)
        s = Statevector(ψ)
        @test s.N == N
        @test s.state == ψ
    end

    @testset "test statevector indexing" begin
        N = rand(1:10)
        ψ = randn(ComplexF64, 2^N)
        s = Statevector(ψ)
        r = rand(1:2^N)
        @test s[r] == ψ[r]
    end

    @testset "test statevector attributes" begin
        N = rand(1:10)
        ψ = randn(ComplexF64, 2^N)
        s = Statevector(ψ)
        @test size(s) == size(ψ)
        @test length(s) == length(ψ)
        @test IndexStyle(typeof(s)) == IndexLinear()
    end

    @testset ExtendedTestSet "util reduced density matrix for statevector" begin
        N = 4
        ψ = Statevector(randn(ComplexF64, 2^N))
        χ = Statevector(randn(ComplexF64, 2^N))
        # full density matrix |ψ⟩⟨χ|
        ρ = reshape(kron(conj(χ.state), ψ.state), 2^N, 2^N)
        id = Matrix{ComplexF64}(I, 2, 2)
        A = randn(ComplexF64, 2, 2)
        B = randn(ComplexF64, 2, 2)
        @testset ExtendedTestSet "reduced density matrix correctness" begin
            @test sum(kron(A, B) .* rdm(N, (4, 2), ψ, χ)) ≈ sum(kron(B, id, A, id) .* ρ)
        end

        @testset ExtendedTestSet "reduced density matrix exceptions" begin
            @test_throws ErrorException("Need at least one wire to act on.") rdm(N, (), ψ, χ)
            @test_throws ErrorException("Number of gate wires cannot be larger than total number of wires.") rdm(N, (1, 2, 3, 4, 5), ψ, χ)
            @test_throws ErrorException("Wire indices must be unique.") rdm(N, (2, 2), ψ, χ)
            @test_throws ErrorException("Wire index cannot be smaller than 1.") rdm(N, (-1, 2), ψ, χ)
            @test_throws ErrorException("Wire index cannot be larger than total number of wires.") rdm(N, (5, 1), ψ, χ)
        end
    end

    @testset ExtendedTestSet "util reduced density matrix for statevector" begin
        N = 4
        ψ = Statevector(randn(ComplexF64, 2^N))
        χ = Statevector(randn(ComplexF64, 2^N))
        # full density matrix |ψ⟩⟨χ|
        ρ = reshape(kron(conj(χ.state), ψ.state), 2^N, 2^N)
        id = Matrix{ComplexF64}(I, 2, 2)
        A = randn(ComplexF64, 2, 2)
        B = randn(ComplexF64, 2, 2)
        @testset ExtendedTestSet "reduced density matrix correctness" begin
            @test sum(kron(A, B) .* rdm(N, (4, 2), ψ, χ)) ≈ sum(kron(B, id, A, id) .* ρ)
        end

        @testset ExtendedTestSet "reduced density matrix exceptions" begin
            @test_throws ErrorException("Need at least one wire to act on.") rdm(N, (), ψ, χ)
            @test_throws ErrorException("Number of gate wires cannot be larger than total number of wires.") rdm(N, (1, 2, 3, 4, 5), ψ, χ)
            @test_throws ErrorException("Wire indices must be unique.") rdm(N, (2, 2), ψ, χ)
            @test_throws ErrorException("Wire index cannot be smaller than 1.") rdm(N, (-1, 2), ψ, χ)
            @test_throws ErrorException("Wire index cannot be larger than total number of wires.") rdm(N, (5, 1), ψ, χ)
        end
    end

end


text:
### Process examples
# Always rerun examples
const EXAMPLES_OUT = joinpath(@__DIR__, "src", "examples")
ispath(EXAMPLES_OUT) && rm(EXAMPLES_OUT; recursive=true)
mkpath(EXAMPLES_OUT)

# Install and precompile all packages
# Workaround for https://github.com/JuliaLang/Pkg.jl/issues/2219
examples = filter!(isdir, readdir(joinpath(@__DIR__, "..", "examples"); join=true))
let script = "using Pkg; Pkg.activate(ARGS[1]); Pkg.instantiate()"
    for example in examples
        if !success(`$(Base.julia_cmd()) -e $script $example`)
            error(
                "project environment of example ",
                basename(example),
                " could not be instantiated",
            )
        end
    end
end
# Run examples asynchronously
processes = let literatejl = joinpath(@__DIR__, "literate.jl")
    map(examples) do example
        return run(
            pipeline(
                `$(Base.julia_cmd()) $literatejl $(basename(example)) $EXAMPLES_OUT`;
                stdin=devnull,
                stdout=devnull,
                stderr=stderr,
            );
            wait=false,
        )::Base.Process
    end
end

# Check that all examples were run successfully
isempty(processes) || success(processes) || error("some examples were not run successfully")

### Build documentation
using Documenter

using KernelFunctions
using PDMats, Kronecker  # we have to load all optional packages to generate the full API documentation

# Doctest setup
DocMeta.setdocmeta!(
    KernelFunctions,
    :DocTestSetup,
    quote
        using KernelFunctions
    end;  # we have to load all packages used (implicitly) within jldoctest blocks in the API docstrings
    recursive=true,
)

makedocs(;
    sitename="KernelFunctions.jl",
    format=Documenter.HTML(),
    modules=[KernelFunctions],
    pages=[
        "Home" => "index.md",
        "userguide.md",
        "kernels.md",
        "transform.md",
        "metrics.md",
        "create_kernel.md",
        "API" => "api.md",
        "Design" => "design.md",
        "Examples" =>
            map(filter!(filename -> endswith(filename, ".md"), readdir(EXAMPLES_OUT))) do x
                return joinpath("examples", x)
            end,
    ],
    strict=true,
    checkdocs=:exports,
    doctestfilters=[
        r"{([a-zA-Z0-9]+,\s?)+[a-zA-Z0-9]+}",
        r"(Array{[a-zA-Z0-9]+,\s?1}|Vector{[a-zA-Z0-9]+})",
        r"(Array{[a-zA-Z0-9]+,\s?2}|Matrix{[a-zA-Z0-9]+})",
    ],
)

deploydocs(;
    repo="github.com/JuliaGaussianProcesses/KernelFunctions.jl.git", push_preview=true
)


text:
struct SolverState{cdim, Solver, Scaling}
	scaling::Scaling
	solver::Solver
	initm::Matrix{Float64}
	initv::Vector{Float64}
	idel::Vector{Float64}
	dx::Vector{Float64}
	dy::Vector{Float64}
	dz::Vector{Float64}
	ds::Vector{Float64}
	nt1::Vector{Float64}
	nt2::Vector{Float64}
	mt1::Vector{Float64}
	kt1::Vector{Float64}
	kt2::Vector{Float64}
	kt3::Vector{Float64}
	rx::Vector{Float64}
	ry::Vector{Float64}
	rz::Vector{Float64}
	rs::Vector{Float64}

	function SolverState(pr::Problem, solver::K) where K <: KKTSolver{S} where S <: AbstractScaling
		scaling = S(pr)
		n = pr.n
		m = pr.m
		k = pr.k

	    initm = zeros(n+m+k, n+m+k)
	    initv = zeros(n+m+k)
		idel = zeros(k)
		dx,dy,dz,ds = zeros(n),zeros(m),zeros(k),zeros(k)
		nt1,nt2 = zeros(n),zeros(n)
		mt1 = zeros(m)
		kt1,kt2,kt3 = zeros(k),zeros(k),zeros(k)
		rx,ry,rz,rs = zeros(n), zeros(m), zeros(k), zeros(k)
		return new{pr.n, K, S}(scaling, solver, initm, initv, idel, dx, dy, dz, ds, nt1, nt2, mt1, kt1, kt2, kt3, rx, ry, rz, rs)
	end
end

function solve_socp(prob::Problem{C,n,m,k,sing}, ss::SolverState{dim, K, S}) where {C,K,S,n,m,k,sing,dim}
	cones = prob.cones

	initm, initv, idel = ss.initm, ss.initv, ss.idel
	dx,dy,dz,ds = ss.dx,ss.dy,ss.dz,ss.ds
	nt1,nt2,mt1,kt1,kt2,kt3 = ss.nt1,ss.nt2,ss.mt1,ss.kt1,ss.kt2,ss.kt3
	rx,ry,rz,rs = ss.rx,ss.ry,ss.rz,ss.rs
	println("$n $m $k")
	#=
    open("A.txt", "w") do f
    	println(f, prob.A)
    end
    open("G.txt", "w") do f
    	println(f, prob.G)
    end
    open("c.txt", "w") do f
    	println(f, prob.c)
    end
    open("b.txt", "w") do f
    	println(f, prob.b)
    end
    open("h.txt", "w") do f
    	println(f, prob.h)
    end
    open("cones.txt", "w") do f
    	println(f, prob.cones)
    end
	=#
	r1 = hcat(spzeros(n,n), prob.A', prob.G')
	r2 = hcat(prob.A, spzeros(m,m), spzeros(m, k))
	r3 = hcat(prob.G, spzeros(k,m), sparse(-I, k, k))
	idmat = vcat(r1, r2, r3)
    initv[1:n] .= (-).(prob.c)
    initv[n+1:n+m] .= prob.b
    initv[n+m+1:n+m+k] .= prob.h
    #=
    open("initm.txt", "w") do f
    	println(f, initm)
    end
    open("initv.txt", "w") do f
    	println(f, initv)
    end
    =#

	initials::Vector{Float64} = idmat\initv

	make_e!(cones, idel)
	iz::Vector{Float64} = initials[n+m+1:n+m+k]
	alphp = max_step(cones, -iz)
	alphd = max_step(cones, iz)

	if abs(alphp) < 1e-10
		inits = -iz
	else
		inits = -iz + (1+alphp)*idel
	end

	if abs(alphd) < 1e-10
		initz = iz
	else
		initz = iz + (1+alphd)*idel
	end

#	println("$initials $inits $initz")
	state = State(prob, initials[1:n], initials[n+1:n+m], initz, inits)
	for i=1:40
		scaling = compute_scaling(cones, ss.scaling, state.s, state.z)
		l = scaling.l
		# solve affine direction
		# dx = prob.A'*state.y .+ prob.G'*state.z .+ prob.c
		mul!(nt1, prob.A', state.y)
		mul!(nt2, prob.G', state.z)
		dx .= nt1 .+ nt2 .+ prob.c
		# dy = prob.A*state.x .- prob.b
		mul!(mt1, prob.A, state.x)
		dy .= mt1 .- prob.b
		# dz = prob.G*state.x .+ state.s .- prob.h
		mul!(kt1, prob.G, state.x)
		dz .= kt1 .+ state.s .- prob.h

		vprod!(cones, ds, l, l)

		if (norm(dx) + norm(dy) + dot(state.z,state.s)) < 1e-5
			break
		end
		rmul!(dx, -1.0); rmul!(dy, -1.0); rmul!(dz, -1.0); rmul!(ds, -1.0)
		setup_iter(ss.solver, prob, state, scaling)
		solve_kkt(ss.solver, prob, state, scaling, dx, dy, dz, ds, rx, ry, rz, rs)
		scale!(prob.cones, scaling, rz, kt3)
		iscale!(prob.cones, scaling, rs, kt2)
		t = compute_step(cones, l, kt3, kt2)

		rho = 1-t-t^2 * dot(kt2, kt3)/dot(l,l)
		sig = max(0, min(1, rho))^3
		mu = dot(l,l)/deg(cones)

		scfact = 1.0-sig
		vprod!(cones, kt1, kt2, kt3)
		mul!(kt2, sig*mu, idel)
		ds .+= kt2 .- kt1
		rmul!(dx, scfact); rmul!(dy, scfact); rmul!(dz, scfact)
		solve_kkt(ss.solver, prob, state, scaling, dx, dy, dz, ds, rx, ry, rz, rs)

		scale!(prob.cones, scaling, rz, kt3)
		iscale!(prob.cones, scaling, rs, kt2)
		step = compute_step(cones, l, kt3, kt2)
		step *= 0.99
		state.x .+= rx .* step
		state.y .+= ry .* step
		state.z .+= rz .* step
		state.s .+= rs .* step
    end
    return state
end


text:
#=
Copyright (c) 2018-2022 Chris Coey, Lea Kapelevich, and contributors

This Julia package Hypatia.jl is released under the MIT license; see LICENSE
file in the root directory or at https://github.com/chriscoey/Hypatia.jl
=#

doptimaldesign_insts(use_logdet::Bool) = [
    [(q, 2q, 2q, 5, use_logdet, !use_logdet, false)
    for q in vcat(3, 50:50:200, 300:100:1000)] # includes compile run
    ]

insts = OrderedDict()
rootdet_insts = doptimaldesign_insts(false)
all_insts = vcat(rootdet_insts, doptimaldesign_insts(true))
insts["nat"] = (nothing, all_insts)
insts["extEP"] = (:ExpPSD, all_insts)
insts["extSEP"] = (:SOCExpPSD, rootdet_insts)
return (DOptimalDesignJuMP, insts)


text:
@testset "test delong" begin
    target = [0,0,1,1,1]
    modela = [0.1, 0.2, 0.6, 0.7, 0.8]
    modelb = [0.3,0.6,0.2,0.7,0.9]
    nontara, tara = modela[1:2], modela[3:end]
    nontarb, tarb = modelb[1:2], modelb[3:end]

    @test round(ROCAnalysis.delong_test(tara, nontara, tarb, nontarb), digits=4) == 0.3173

    target = [0,0,0,0,0,0,1,1,1,1,1,1,1]
    modela = [0.1,0.2,0.05,0.3,0.1,0.6,0.6,0.7,0.8,0.99,0.8,0.67,0.5]
    modelb = [0.3,0.6,0.2,0.1,0.1,0.9,0.23,0.7,0.9,0.4,0.77,0.3,0.89]
    nontara, tara = modela[1:6], modela[7:end]
    nontarb, tarb = modelb[1:6], modelb[7:end]

    @test round(ROCAnalysis.delong_test(tara, nontara, tarb, nontarb), digits=5) == 0.09453
end


text:
println("Running runtests.jl ...")
push!(LOAD_PATH, "..")

using FunctionalData
using HDF5
using MultivariateStats
using ShapeModels
using Statistics
using Test

macro shouldtestset(a,b) length(ARGS) < 1 || ARGS[1] == a ?  :(@testset $a $b) : nothing end
macro shouldtestset2(a,b) length(ARGS) < 2 || ARGS[2] == a ?  :(@testset $a $b) : nothing end

landmarks = ShapeModels.examplelandmarks()

@shouldtestset "basic" begin
	a = PCAShapeModel(landmarks)
	@test size(projection(a.pca)) == (256,8)
	@test size(principalvars(a.pca)) == (8,)
	@test indim(a.pca) == 256
	@test outdim(a.pca) == 8
end

@shouldtestset "utils" begin
	a = PCAShapeModel(landmarks)
	m = meanshape(a)
	@test size(m) == (2,128)
	@test maximum(abs.(mean(m, dims = 2)))  < 1e-6
end

@shouldtestset "shape" begin
	a = PCAShapeModel(landmarks)

    @test size(shape(a, zeros(nmodes(a)))) == (2,128)
end

@shouldtestset "modes" begin
	a = PCAShapeModel(landmarks)
	@test size(modeshapes(a, 1)) == (2,128,10)
end

println("  ... finished runtests.jl")


text:
"""
    SubsampledDataCausalityTest

A causality test where the data is subsampled in some manner before
applying the test.
"""
abstract type SubsampledDataCausalityTest{CT} <: MetaCausalityTest{CT} end

get_ηs(x::SubsampledDataCausalityTest) = get_ηs(x.test)

export SubsampledDataCausalityTest, get_ηs

include("RandomSequencesTest.jl")
include("causality_RandomSequencesTest.jl")


text:
module libvulkan
	using CBinding
	using Vulkan_Headers_jll

	c`$([
		Main.SYSROOT...,

		"-I$(Vulkan_Headers_jll.artifact_dir)/include",
		"-DVK_NO_PROTOTYPES",
	])`

	const c"int8_t" = Int8
	const c"int16_t" = Int16
	const c"int32_t" = Int32
	const c"int64_t" = Int64
	const c"uint8_t" = UInt8
	const c"uint16_t" = UInt16
	const c"uint32_t" = UInt32
	const c"uint64_t" = UInt64

	const c"size_t" = Csize_t

	c"""
		#include <vulkan/vulkan.h>
	"""ji
end


text:
using IterativeSolvers
using ExtendableSparse
using LinearAlgebra

mutable struct Wendland{X,Y,L,U,C,I,T,E} <: AbstractSurrogate
    x::X
    y::Y
    lb::L
    ub::U
    coeff::C
    maxiters::I
    tol::T
    eps::E
end

@inline _l(s,k) = floor(s/2) + k + 1

function _wendland(x,eps)
    r = eps * norm(x)
    val = (1.0 - r)
    if val >= 0
        dim = length(x)
        #at the moment only k = 1 is supported, but we could also support
        # missing wendland (k=1/2,k=3/2,k=5/2), and different k's.
        ell = _l(dim,1)
        powerTerm = ell + 1.0
        val = (1.0 - r)
        return val^powerTerm * (powerTerm * r + 1.0)
    else
        return zero(eltype(x[1]))
    end
end

function _calc_coeffs_wend(x,y,eps,maxiters,tol)
    n = length(x)
    W = ExtendableSparseMatrix{eltype(x[1]),Int}(n,n)
    @inbounds for i = 1:n
        k = i #wendland is symmetric
        for j = k:n
            W[i,j] = _wendland(x[i] .- x[j],eps)
        end
    end
    U = Symmetric(W,:U)
    return cg(U,y,maxiter = maxiters, reltol = tol)
end

function Wendland(x,y,lb,ub; eps = 1.0, maxiters = 300, tol = 1e-6)
    c = _calc_coeffs_wend(x,y,eps,maxiters,tol)
    return Wendland(x,y,lb,ub,c,maxiters,tol,eps)
end

function (wend::Wendland)(val)
    return sum(wend.coeff[j]*_wendland(val,wend.eps) for j=1:length(wend.coeff))
end

function add_point!(wend::Wendland,new_x,new_y)
    if (length(new_x) == 1 && length(new_x[1]) == 1) || ( length(new_x) > 1 && length(new_x[1]) == 1 && length(wend.lb)>1)
        push!(wend.x,new_x)
        push!(wend.y,new_y)
    else
        append!(wend.x,new_x)
        append!(wend.y,new_y)
    end
    wend.coeff = _calc_coeffs_wend(wend.x,wend.y,wend.eps,wend.maxiters,wend.tol)
    nothing
end


text:
using Bridge, Random, Test, StaticArrays

struct Diff_Lyap{N} <: ContinuousTimeProcess{SVector{N,Float64}}
end
Diff_Lyap(d) = Diff_Lyap{d}()

dim(::Diff_Lyap{N}) where {N} = N
@testset "Lyap" begin
    # example setting
    d = 20
    t = 2.0:0.01:5.0

    Random.seed!(4)
    Bridge.B(t, P::Diff_Lyap) = reshape(rand(dim(P)^2),dim(P),dim(P))
    Bridge.σ(t, P::Diff_Lyap) = reshape(rand(10*dim(P)),dim(P),10)
    Bridge.a(t, P::Diff_Lyap) = Bridge.σ(t,P) * Bridge.σ(t,  P)'
    Pt = Diff_Lyap(d)

    Hend = Matrix(1.0I, d,d)
    out = Bridge.lyapunovpsdbackward(t, Pt, Hend)

    @test all(map(x-> isposdef(Matrix(Bridge.symmetrize(x))), out.yy))
end


text:
### A Pluto.jl notebook ###
# v0.14.2

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 9a0cec14-08db-11eb-3cfa-4d1c327c63f1
begin
    import Pkg

    Pkg.activate(mktempdir())
    Pkg.add([
        Pkg.PackageSpec(name="Plots", version="1"),
        Pkg.PackageSpec(name="PlutoUI", version="0.7"),
        Pkg.PackageSpec(name="StatsBase", version="0.33"),
    ])

    using Plots, PlutoUI, StatsBase, Statistics
end

# ╔═╡ 41f7d874-8cb9-11eb-308d-47dea998f6bf
html"""

Section 3.1

 Time stepping

Lecture Video

"""

# ╔═╡ fb6cdc08-8b44-11eb-09f5-43c167aa53fd
PlutoUI.TableOfContents(aside=true)

# ╔═╡ 6f871ac0-716d-4bf8-a067-c798869c103f
md"""
# Modeling component failure: Discrete and continuous
"""

# ╔═╡ ae243395-521b-4834-b61e-19501e54b41c
md"""
Let's think about a simple model for failure of components such as light bulbs.
(We actually derived this global, or macroscopic, model from a microscopic stochastic model a few lectures ago.)

Components can fail at any moment but we'll start off by checking once per day to count the number that have failed during that day. Then we'll check several times per day. This is still a **discrete** model. Finally we'll see how we can turn this into a **continuous** model where we can talk about the number that have failed by any real time $t$.
"""

# ╔═╡ 8d2858a4-8c38-11eb-0b3b-61a913eed928
md"""
## Checking failures once per day (integer time steps)
"""

# ╔═╡ e93c5f2f-d7c8-41ea-bdbb-7cf6587b6266
md"""

Let's call $N_k$ the (average) number of bulbs that are still functioning on day number $k$, starting from an initial number $N_0$.

We can find an equation for the number $N_{k+1}$ that are still functioning at the end of day number $k+1$ by working out how many *fail* on day $k+1$.

Let's call $p$ the probability that each bulb fails each day. For example, if 10% of the bulbs fail each day then $p = 0.1$.

If there are 100 bulbs and 10% fail on the first day, then 10 fail, so there will be 90 remaining. In general, if a proportion $p$ of the $N_k$ fail, in total we expect $p \, N_k$ to fail.

Hence

$$N_{k+1} = N_k - p \, N_k$$

or

$$N_{k+1} - N_k = - p N_k.$$

In this very simple model we can actually solve this recurrence relation analytically to find the number still functioning at time $t$:

$$N_{k+1} = (1 - p) N_k$$

so

$$N_k = (1 - p) N_{k-1} = (1 - p)^2 N_{k-2} = \cdots$$

hence

$$N_k = N_0 \, (1 - p)^k.$$

"""

# ╔═╡ 43f5ac88-7d07-429c-b27f-49908c30bdf9
md"""
## Checking failures $n$ times per day
"""

# ╔═╡ b70465ab-9c7c-4533-9539-b414ef54a892
md"""
Now suppose instead we ask how many light bulbs fail in *half* a day. If 10% fail per day, it's natural to think that 5% fail in half a day.

However, this is not quite right due to the effect of **compounding** (as in compound interest). If 5% fail and then 5% of the remainder fail, then how many are left?

"""

# ╔═╡ 80c1a728-1784-4bb7-b2c3-e77b41929a78
(1 - 0.05) * (1 - 0.05)

# ╔═╡ a25cb8e6-f65d-4061-b90c-079814458c94
md"""
So slightly *fewer* than 10% in total have failed, due to the effect of compounding. Nonetheless, the result is approximately right, so let's take that.
"""

# ╔═╡ 4cc483e5-c322-44c8-83f1-802b6cb432aa
md"""
So let's suppose that $n$ times a day, a proportion $p / n$ fail (for example, 10% and twice a day gives 5% failing each time, approximately).

Then the number remaining after the first failure check on day $k$ is
"""

# ╔═╡ 3ce501b4-76bc-49ab-b3b8-a41f29dbcc2b
md"""
$N_{k + \frac{1}{n}} = \textstyle (1 - \frac{p}{n}) N_k$

Here we have used a subscript since we are in a discrete situation. We could also have written instead

$$N(k + \textstyle \frac{1}{n}).$$

The solution at the next day is

$$N_{k+1} = N_k \, (1 - \textstyle \frac{p}{n})^n.$$

And the full solution for the number remaining after $k$ days is

$N_k = N_0 \, \textstyle (1 - \frac{p}{n})^{nk}$
"""

# ╔═╡ c539e622-d76d-489a-abb9-4ba47dfe9b90
md"""
Let's plot these to see what they look like:
"""

# ╔═╡ 95fe38c5-d717-47d0-8db7-5f8d53a6c6f1
md"""
n = $(n_slider = @bind n Slider(0:8, show_value=true))
"""

# ╔═╡ 2982c418-dad5-44cc-8194-5b607af84b16
p = 0.4

# ╔═╡ c97964d1-b5d2-4ee7-80cc-995b3f344aa1
let
	N0 = 100
	T = 20

	N = [N0 * (1 - p)^t for t in 0:T]

	plot(0:T, N, m=:o, alpha=0.5, ms=3, label="once daily", lw=2)

# 	N2 = [N0 * (1 - p5/2)^(t) for t in 0:2T]
# 	plot!(0:0.5:T, N2, m=:o, alpha=0.5, ms=2, label="twice a day")

# 	N4 = [N0 * (1 - p5/4)^(t) for t in 0:4T]
# 	plot!(0:0.25:T, N4, m=:o, alpha=0.5, ms=2, label="four times a day")

	N = [N0 * (1 - p/(2^n))^(t) for t in 0:(2^n)*T]
	plot!(0:(2.0^(-n)):T, N, m=:o, alpha=0.5, ms=2, label="$(2^n) times per day")

	xlabel!("days (k)")
	ylabel!("N_k")

	title!("$(2^n) times per day")

	# plot!(t -> N0 * exp(-p5 * t))
end

# ╔═╡ ba121b40-2bfc-42d4-81ee-5f90e18ec8de
md"""
## Continuous time
"""

# ╔═╡ 74892ec6-6639-469d-8711-5039a140d833
md"""
Thinking back to the class on "discrete to continuous", we see that we are producing more and more discrete values to keep track of, but after a while the curve they trace out does not really change. It thus makes sense to define a **limiting object**, which is what would happen if you could take smaller and smaller time steps in the right way, i.e. take the limit as $n \to \infty$.

In that case we could imagine being able to calculate the (average) number $N(t)$
of bulbs that are functioning at time $t$, where $t$ can be *any* positive real number.

"""

# ╔═╡ fc6899d3-ea18-487a-add1-20be86ce9c74
md"""
In calculus we learn ways to see that
"""

# ╔═╡ 75a60bcf-3f77-49fb-a7ee-db4580aae6f3
md"""
$(1 - p/n)^n$ converges to $\exp(-p)$ as $n \to \infty$.
"""

# ╔═╡ 786beb46-d175-44b3-a63e-057150e53c66
md"""
An alternative approach is to look at the time evolution in terms of *differences*. In a time $1/n$ a proportion $p/n$ decays. Here we are thinking of $\delta t = 1/n$ as the **time step** between consecutive checks. If there are $N(t)$ bulbs functioning at a time $t$ and we take a small time  step of general length $\delta t$, a proportion $p \, \delta t$ should fail, so that
"""

# ╔═╡ 2d9f3aad-1d41-4918-a128-47dc58b667e3
md"""
$N(t + \delta t) - N(t) = -(p \, \delta t) \, N(t).$
"""

# ╔═╡ 17a23094-3975-49f4-8c7f-03fbc8afbbf2
md"""
Dividing through by $\delta t$ we find
"""

# ╔═╡ eaf6f4eb-b367-492e-a1be-81f9455252c4
md"""
$$\frac{N(t + \delta t) - N(t)}{\delta t} = -p \, N(t).$$

"""

# ╔═╡ 51c226b9-ab3d-46b4-a963-3548ad715d85
md"""
We now recognise the left-hand side of the equation: if we take the limit as $\delta t \to 0$, we have exactly the definition of  the **derivative** $\frac{dN(t)}{dt}$.

Hence, taking that limit we obtain
"""

# ╔═╡ 6c527098-ab53-4862-bda6-0c11b1564a11
md"""
$$\frac{dN(t)}{dt} = - p \, N(t)$$

with $N(0) = N_0$, the initial number.

This is an **ordinary differential equation**: it is an equation relating the value of the function $N(t)$ at time $t$ to the derivative (slope) of that function at that point. This relationship must hold for all $t$. It is not obvious that this equation even makes sense (although it should do, given the way we have derived it), but in differential equations courses we see that it does make sense (under some technical conditions), and uniquely specifies a function satisfying the ODE together with the initial condition. (This is called an "initial-value problem".)
"""

# ╔═╡ 3c2b2f03-522c-40a0-ac1d-8054fe8e3fa2
md"""
In this particular case, once again we are lucky enough to be able to solve this equation analytically:  $N(t)$ is a function whose derivative is a multiple of the same function, and hence it must be exponential:
"""

# ╔═╡ 5cec433e-ee71-44b5-b5d6-3feab80fa535
md"""
$N(t) = N_0 \exp(-p \, t)$
"""

# ╔═╡ 96b02ce7-ce16-4276-a147-ba94d7a2e160
md"""
This is an alternative way to define the exponential function. We can add this to the above plot:
"""

# ╔═╡ d952db33-1f82-42f5-96af-8038c256715b
n_slider

# ╔═╡ 2b9276dc-fcca-4469-a62c-028a9eb3c2a9
let
	N0 = 100
	T = 20

	N = [N0 * (1 - p)^t for t in 0:T]

	plot(0:T, N, m=:o, alpha=0.5, ms=3, label="once daily", lw=2)

# 	N2 = [N0 * (1 - p5/2)^(t) for t in 0:2T]
# 	plot!(0:0.5:T, N2, m=:o, alpha=0.5, ms=2, label="twice a day")

# 	N4 = [N0 * (1 - p5/4)^(t) for t in 0:4T]
# 	plot!(0:0.25:T, N4, m=:o, alpha=0.5, ms=2, label="four times a day")

	N = [N0 * (1 - p/(2^n))^(t) for t in 0:(2^n)*T]
	plot!(0:(2.0^(-n)):T, N, m=:o, alpha=0.5, ms=2, label="$(2^n) times per day")

	xlabel!("days (k)")
	ylabel!("N_k")

	title!("$(2^n) times per day")

	plot!(t -> N0 * exp(-p * t), label="continuous", lw=2)
end

# ╔═╡ 754fe8c1-7021-48e8-9523-d5b22d0af93f
md"""
We see graphically that this is indeed the correct limiting curve.
"""

# ╔═╡ ccb35ad7-db20-46fa-abff-a6e88ef999e0
md"""
In this context, $p$ is called a **rate**; it is a *probability per unit time*, i.e. a ratio of probability and time. (To get the probability of decaying in a time $\delta t$, we *multiplied* $p$ by $\delta t$.)
"""

# ╔═╡ d03d9bfc-20ea-49bc-bc7b-df22cc240ffe
md"""
Let's summarise what we have found:
"""

# ╔═╡ dde3ffc7-b333-4305-b8e7-9888e4512c41
md"""
| Step type   | Time stepping     |  Difference    |   Solution  |
| ----------- | :-----------: | :-----------:  | :-----------: |
| Integer            | $N_{k+1} = (1 - p) N_k$  |  $N_{k+1} - N_k = -p N_k$  | $N_k = N_0 (1 - p)^k$
| Rational |  $N_{k + \frac{1}{n}} = \textstyle (1 - \frac{p}{n}) N_k$ | $N_{k + \frac{1}{n}} - N_k = \textstyle (- \frac{p}{n}) N_k$   | $N_k = N_0 (1 - \frac{p}{n})^{n k}$
| Continuous   | $N(t + \delta t) = (1 - p \, \delta t) N(t)$         |  $\frac{dN(t)}{dt} = - p \, N(t)$  |  $N(t) = N_0 \exp(-p \, t)$

"""

# ╔═╡ d74bace6-08f4-11eb-2a6b-891e52952f57
md"""
# SIR model
"""

# ╔═╡ 76268535-e232-4e02-97cd-cf9b3ddec256
md"""

Let's look at a more complicated example, the **SIR** model of the spread of an epidemic, or of a rumour, in a population. You are surely familiar with models of this type modelling the spread of COVID-19, and from the homework.

As in the homework, we can make a fully discrete, stochastic agent-based model, where we give microscopic rules specifying how individual agents interact with one another. When we run such models with large enough systems, we observe that the results are often quite smooth. An alternative approach is to try to write down macroscopic discrete equations that describe the dynamics of averages.

Often it is easier to understand the behaviour of such systems by formulating a continuous version of the model. Some people make discrete models because they're not happy with continuous models; large discrete models can also be computationally wasteful. On the other hand, they can also include effects that might be more difficult to model in a continuous framework, e.g. non-local effects or things that "don't become continuous very well".
"""

# ╔═╡ 11e24e1d-39db-4b7e-96db-50458def72af
md"""
## Discrete-time SIR model
"""

# ╔═╡ dbdf2812-08f4-11eb-25e7-811522b24627
md"""
First let's think about the SI model: agents can be susceptible (S) and infectious (I). A susceptible person becomes infectious when they come into contact with an infectious person, with some probability.
"""

# ╔═╡ 238f0716-0903-11eb-1595-df71600f5de7
md"""
Let's call $S_t$ and $I_t$ be the number of susceptible and infectious people at time $t$, respectively, and let's call $N$ the total number of people.

Let's suppose that at each time step, each infectious person has the chance to interact with one other person (on average). That person will be chosen at random from the total population of size $N$. A new infection occurs only if that chosen person is susceptible, which happens with probability $S_t / N$, and only if the infection attempt is successful, with probability $b$, say.

Hence the change in the number of infectious people after that step is
"""

# ╔═╡ 8e771c8a-0903-11eb-1e34-39de4f45412b
md"""
$$\Delta I_t = I_{t+1} - I_t = b \, I_t \, \left(\frac{S_t}{N} \right)$$
"""

# ╔═╡ fb52c62d-15d3-46a2-8e3d-2de20c68ded4
md"""
The decrease in $S_t$ is also given by $\Delta I_t$.
"""

# ╔═╡ e83fc5b8-0904-11eb-096b-8da3a1acba12
md"""
There is also recovery, with a constant probability $c$ at each step, once you are infectious.

It is useful to normalize by $N$, so we define the proportions of the population that are susceptible, infectious and recovered as

$$s_t := \frac{S_t}{N}; \quad i_t := \frac{I_t}{N}; \quad r_t := \frac{R_t}{N}.$$
"""

# ╔═╡ d1fbea7a-0904-11eb-377d-690d7a16aa7b
md"""
Including recovery with probability $c$ we obtain the **discrete-time SIR model**:
"""

# ╔═╡ dba896a4-0904-11eb-3c47-cbbf6c01e830
md"""
$$\begin{align}
s_{t+1} &= s_t - b \, s_t \, i_t \\
i_{t+1} &= i_t + b \, s_t \, i_t - c \, i_t\\
r_{t+1} &= r_t + c \, i_t
\end{align}$$
"""

# ╔═╡ cea2dcfb-b1eb-4269-81d7-8596969e9bd6
md"""
## Continuous-time SIR model
"""

# ╔═╡ 08d166f1-3af0-45a8-bcad-6ee958497453
md"""
We can now go through the same process as with the failure model, where we take time steps of length $\delta t$ instead, and replace probabilities $b$ and $c$ with *rates* $\beta$ and $\gamma$. Taking the limit $\delta t \to 0$ we get
"""

# ╔═╡ 72061c66-090d-11eb-14c0-df619958e2b6
md"""
$$\begin{align}
\frac{ds(t)}{dt} &= -\beta \, s(t) \, i(t) \\
\frac{di(t)}{dt} &= +\beta \, s(t) \, i(t) &- \gamma \, i(t)\\
\frac{dr(t)}{dt} &= &+ \gamma \, i(t)
\end{align}$$
"""

# ╔═╡ c07367be-0987-11eb-0680-0bebd894e1be
md"""
We can think of this as a model of a chemical reaction with species S, I and R. The term $s(t) i(t)$ is known as the [**mass action**](https://en.wikipedia.org/wiki/Law_of_mass_action) form of interaction.

Note that no analytical solutions of these (simple) nonlinear ODEs are known as a function of time! (However, [parametric solutions are known](https://arxiv.org/abs/1403.2160).)
"""

# ╔═╡ f8a28ba0-0915-11eb-12d1-336f291e1d84
md"""
Below is an example simulation of the discrete-time model.
"""

# ╔═╡ d994e972-090d-11eb-1b77-6d5ddb5daeab
begin
	NN = 100

	SS = NN - 1
	II = 1
	RR = 0
end

# ╔═╡ 050bffbc-0915-11eb-2925-ad11b3f67030
ss, ii, rr = SS/NN, II/NN, RR/NN

# ╔═╡ 1d0baf98-0915-11eb-2f1e-8176d14c06ad
p_infection, p_recovery = 0.1, 0.01

# ╔═╡ 28e1ec24-0915-11eb-228c-4daf9abe189b
TT = 1000

# ╔═╡ 349eb1b6-0915-11eb-36e3-1b9459c38a95
function discrete_SIR(s0, i0, r0, T=1000)

	s, i, r = s0, i0, r0

	results = [(s=s, i=i, r=r)]

	for t in 1:T

		Δi = p_infection * s * i
		Δr = p_recovery * i

		s_new = s - Δi
		i_new = i + Δi - Δr
		r_new = r      + Δr

		push!(results, (s=s_new, i=i_new, r=r_new))

		s, i, r = s_new, i_new, r_new
	end

	return results
end

# ╔═╡ 39c24ef0-0915-11eb-1a0e-c56f7dd01235
SIR = discrete_SIR(ss, ii, rr)

# ╔═╡ 442035a6-0915-11eb-21de-e11cf950f230
begin
	ts = 1:length(SIR)
	discrete_time_SIR_plot = plot(ts, [x.s for x in SIR],
		m=:o, label="S", alpha=0.2, linecolor=:blue, leg=:right, size=(400, 300))
	plot!(ts, [x.i for x in SIR], m=:o, label="I", alpha=0.2)
	plot!(ts, [x.r for x in SIR], m=:o, label="R", alpha=0.2)

	xlims!(0, 500)
end

# ╔═╡ 5ce67dae-96a4-4264-b885-d9f894459330
md"""
Try to implement a model for pendulum using the discreet method and time stepping and convert to continuous domain solution. What does it mean when Dave Sanders said that these equations do not converse energy which is necessary for modelling a pendulum.
"""

# ╔═╡ 5f5d7332-b5f8-4d05-971b-ec0564f1339b
md"""
# Time stepping: The Euler method
"""

# ╔═╡ 7cf51986-5983-4094-a18f-f95f2f6993da
md"""
Above we showed how we can think of Ordinary Differential Equations (ODEs) as arising in a natural way from discrete-time models where we take time steps.

What about if somebody gives us an ODE -- how should we solve it numerically?

In fact, we do the reverse process: We **discretize** the equation and reduce it to a system where we take time steps!

Suppose the differential equation is

$$\dot{x} = f(x)$$

The simplest such method is the **Euler method**: we approximate the derivative using a small (but not *too* small) time step $h$ (what we called $\delta t$ above):

$$\dot{x} \simeq \frac{x(t + h) - x(t)}{h}$$

giving

$$x(t + \delta t) \simeq x(t) + h \, f(x).$$

The Euler method with constant time step is then the following algorithm:

$$x_{k+1} = x_k + h \, f(x_k)$$

If we have several variables in our ODE, we can wrap the variables up into a vector and use the *same* method: for the ODE

$$\dot{\mathbf{x}} = \mathbf{f}(\mathbf{x})$$

the Euler method becomes

$$\mathbf{x}_{k+1} = \mathbf{x}_k + h \, \mathbf{f}(\mathbf{x}_k),$$

where $\mathbf{f}$ denotes is a vector-valued function mapping the vector of variables to the vector of right-hand sides of the ODEs.

"""

# ╔═╡ 763bbb15-c52e-4159-99b7-f3d17f47d56a
md"""
However, in general the Euler method is *not* a good algorithm to simulate the dynamics of an ODE! We can see why that might be from the graphs at the start of this notebook: taking time steps like this actually does a *bad* job at approximating the continuous curve. Numerical analysis courses show how to design better numerical methods to approximate the true solution of an ODE more accurately.

The Julia [SciML / DifferentialEquations.jl](https://diffeq.sciml.ai/stable/tutorials/ode_example/) ecosystem provides a large suite of methods for solving ODEs and many other types of differential equations using state-of-the-art methods.
"""

# ╔═╡ Cell order:
# ╟─41f7d874-8cb9-11eb-308d-47dea998f6bf
# ╠═9a0cec14-08db-11eb-3cfa-4d1c327c63f1
# ╠═fb6cdc08-8b44-11eb-09f5-43c167aa53fd
# ╟─6f871ac0-716d-4bf8-a067-c798869c103f
# ╟─ae243395-521b-4834-b61e-19501e54b41c
# ╟─8d2858a4-8c38-11eb-0b3b-61a913eed928
# ╟─e93c5f2f-d7c8-41ea-bdbb-7cf6587b6266
# ╟─43f5ac88-7d07-429c-b27f-49908c30bdf9
# ╟─b70465ab-9c7c-4533-9539-b414ef54a892
# ╠═80c1a728-1784-4bb7-b2c3-e77b41929a78
# ╟─a25cb8e6-f65d-4061-b90c-079814458c94
# ╟─4cc483e5-c322-44c8-83f1-802b6cb432aa
# ╟─3ce501b4-76bc-49ab-b3b8-a41f29dbcc2b
# ╟─c539e622-d76d-489a-abb9-4ba47dfe9b90
# ╟─95fe38c5-d717-47d0-8db7-5f8d53a6c6f1
# ╠═2982c418-dad5-44cc-8194-5b607af84b16
# ╠═c97964d1-b5d2-4ee7-80cc-995b3f344aa1
# ╟─ba121b40-2bfc-42d4-81ee-5f90e18ec8de
# ╟─74892ec6-6639-469d-8711-5039a140d833
# ╟─fc6899d3-ea18-487a-add1-20be86ce9c74
# ╟─75a60bcf-3f77-49fb-a7ee-db4580aae6f3
# ╟─786beb46-d175-44b3-a63e-057150e53c66
# ╟─2d9f3aad-1d41-4918-a128-47dc58b667e3
# ╟─17a23094-3975-49f4-8c7f-03fbc8afbbf2
# ╟─eaf6f4eb-b367-492e-a1be-81f9455252c4
# ╟─51c226b9-ab3d-46b4-a963-3548ad715d85
# ╟─6c527098-ab53-4862-bda6-0c11b1564a11
# ╟─3c2b2f03-522c-40a0-ac1d-8054fe8e3fa2
# ╟─5cec433e-ee71-44b5-b5d6-3feab80fa535
# ╟─96b02ce7-ce16-4276-a147-ba94d7a2e160
# ╠═d952db33-1f82-42f5-96af-8038c256715b
# ╠═2b9276dc-fcca-4469-a62c-028a9eb3c2a9
# ╟─754fe8c1-7021-48e8-9523-d5b22d0af93f
# ╟─ccb35ad7-db20-46fa-abff-a6e88ef999e0
# ╟─d03d9bfc-20ea-49bc-bc7b-df22cc240ffe
# ╟─dde3ffc7-b333-4305-b8e7-9888e4512c41
# ╟─d74bace6-08f4-11eb-2a6b-891e52952f57
# ╟─76268535-e232-4e02-97cd-cf9b3ddec256
# ╟─11e24e1d-39db-4b7e-96db-50458def72af
# ╟─dbdf2812-08f4-11eb-25e7-811522b24627
# ╟─238f0716-0903-11eb-1595-df71600f5de7
# ╟─8e771c8a-0903-11eb-1e34-39de4f45412b
# ╟─fb52c62d-15d3-46a2-8e3d-2de20c68ded4
# ╟─e83fc5b8-0904-11eb-096b-8da3a1acba12
# ╟─d1fbea7a-0904-11eb-377d-690d7a16aa7b
# ╟─dba896a4-0904-11eb-3c47-cbbf6c01e830
# ╟─cea2dcfb-b1eb-4269-81d7-8596969e9bd6
# ╟─08d166f1-3af0-45a8-bcad-6ee958497453
# ╟─72061c66-090d-11eb-14c0-df619958e2b6
# ╟─c07367be-0987-11eb-0680-0bebd894e1be
# ╟─f8a28ba0-0915-11eb-12d1-336f291e1d84
# ╠═442035a6-0915-11eb-21de-e11cf950f230
# ╠═d994e972-090d-11eb-1b77-6d5ddb5daeab
# ╠═050bffbc-0915-11eb-2925-ad11b3f67030
# ╠═1d0baf98-0915-11eb-2f1e-8176d14c06ad
# ╠═28e1ec24-0915-11eb-228c-4daf9abe189b
# ╠═349eb1b6-0915-11eb-36e3-1b9459c38a95
# ╠═39c24ef0-0915-11eb-1a0e-c56f7dd01235
# ╠═5ce67dae-96a4-4264-b885-d9f894459330
# ╟─5f5d7332-b5f8-4d05-971b-ec0564f1339b
# ╟─7cf51986-5983-4094-a18f-f95f2f6993da
# ╟─763bbb15-c52e-4159-99b7-f3d17f47d56a


text:
module Tree

#=
TODO: Re-implement nodistances and diameter. It's WRONG and slow.

A function that calculates if two trees are identical
With and without considering node labels
With and without considering node lengths
    (can be implemented by operating on copied 0-length tree)
=#

struct Pairs{T}
    x::Vector{T}
end

function Base.iterate(x::Pairs, state::Tuple{Int,Int}=(1,2))
    ai, bi = state
    len = length(x.x)
    if bi > len
        ai += 1
        bi = ai + 1
    end
    bi > len && return nothing
    val = @inbounds (x.x[ai], x.x[bi])
    return val, (ai, bi+1)
end

Base.IteratorSize(::Type{Pairs}) = Base.HasLength()
Base.IteratorEltype(::Type{Pairs}) = Base.HasEltype()
Base.eltype(::Type{Pairs{T}}) where T = Tuple{T, T}
Base.length(x::Pairs) = (length(x.x) * (length(x.x)-1)) >>> 1

abstract type Node end

mutable struct Child <: Node
    name::String
    parent::Node
    length::Float64
    children::Array{Child,1}

    function Child(name, parent, length)
        self = new(name, parent, length, Child[])
        push!(parent.children, self)
        return self
    end
end

Child(name::String, parent::Node) = Child(name, parent, 0.0)

mutable struct Root <: Node
    name::String
    children::Array{Child,1}
    function Root(name)
        new(name, Child[])
    end
end

# This funcion only exists because mutually recursive types cannot exist
# in Julia as of 20191004, so I need some way to tell the compiler that
# child.parent returns Union{Child, Root}, not Node
function Base.getproperty(x::Child, sym::Symbol)
    if sym === :parent
        return getfield(x, :parent)::Union{Child, Root}
    else
        return getfield(x, sym)
    end
end

function Base.show(io::IO, node::Root)
    print(io, "*$(node.name)($(join([c.name for c in node.children], ", ")))")
end

function Base.show(io::IO, node::Child)
    print(io, "$(node.name)($(join([c.name for c in node.children], ", ")))")
end

function Base.deepcopy_internal(x::Child, dict::IdDict)
    if haskey(dict, x)
        return dict[x]
    elseif isempty(dict)
        y = Child(x.name, Root("x"), x.length)
        y.parent = x.parent
    else
        y = Child(x.name, dict[x.parent], x.length)
    end

    dict[x] = y

    for child in x.children
        Base.deepcopy_internal(child, dict)
    end

    return y
end

function Base.deepcopy_internal(x::Root, dict::IdDict)
    y = Root(x.name)
    dict[x] = y
    for child in x.children
        Base.deepcopy_internal(child, dict)
    end
    return y
end

function deepcopyasroot(x::Child)
    dict = IdDict()
    y = Root(x.name)
    dict[x] = y
    for child in x.children
        Base.deepcopy_internal(child, dict)
    end

    return y
end

deepcopyasroot(x::Root) = deepcopy(x)

function removechild!(child::Child)
    deleteat!(child.parent.children, findfirst(isequal(child), child.parent.children))
end

function Base.delete!(node::Child)
    "Deletes a node returning it, while connecting the parent and the children of it."

    length = node.length
    parent = node.parent

    index = findfirst(isequal(node), parent.children)
    deleteat!(parent.children, index)

    for child in node.children
        child.length += length
        child.parent = parent

        # Make sure the newly added children are in the same position in the
        # children array as the old child was.
        insert!(parent.children, index, child)
        index += 1
    end

    return node
end

function insert(name, node::Child, length)
    "Inserts a node on a branch at 'length' distance ancestral from 'child'."

    if length > node.length
        error("length is longer than the branch leading to its child")
    end
    newchild = Child(name, node.parent, node.length - length)
    push!(newchild.children, node)
    removechild!(node)
    node.parent = newchild
    node.length = length

    return newchild
end

function detach!(node::Child; asroot=false)
    """Deletes a subtree from a node, returning it.
    If asroot, convert the given node to a root."""

    removechild!(node)
    asroot || return node
    root = Root(node.name)
    root.children = copy(node.children)
    for child in root.children
        child.parent = root
    end
    return root
end

function move!(node::Child, parent::Node)
    "Moves a child to a new parent"
    removechild!(node)
    push!(parent.children, node)
    node.parent = parent
end

function namemapof(node::T; unique=true) where T <: Node
    "Returns a name => node dict of all nodes in subtree"

    namemap = Dict{String, Node}(node.name=>node)
    for child in descendantsof(node)
        if unique && haskey(namemap, child.name)
            error("$(node.name) is present more than once in tree")
        end
        namemap[child.name] = child
    end
    return namemap
end

abstract type AbstractLineageIterator{T} end
Base.IteratorEltype(::Type{<:AbstractLineageIterator}) = Base.HasEltype()
Base.IteratorSize(::Type{<:AbstractLineageIterator}) = Base.SizeUnknown()

struct ChildLineageIterator <: AbstractLineageIterator{Child}
    x::Child
end

Base.iterate(x::ChildLineageIterator, state=x.x) = (state isa Root) ? nothing : (state, state.parent)
Base.eltype(::Type{ChildLineageIterator}) = Child

struct LineageIterator{T} <: AbstractLineageIterator{T}
    x::Union{T}
end
LineageIterator(x::T) where {T<:Node} = LineageIterator{T}(x)

function Base.iterate(x::LineageIterator, state=x.x)
    state === nothing && return nothing
    return (state isa Root) ? (state, nothing) : (state, state.parent)
end
Base.eltype(::Type{LineageIterator{Child}}) = Union{Root, Child}
Base.eltype(::Type{LineageIterator{Root}}) = Root

lineageof(node::Node) = LineageIterator(node)
childlineageof(node::Child) = ChildLineageIterator(node)

rootof(node::Root) = node
function rootof(node::Child)::Root
    for i in lineageof(node)
        isa(i, Root) && return i
    end
end

siblingsof(node::Child)::Vector{Child} = node.parent.children

struct DepthsFirst{T}
    remaining::Vector{Union{Child, T}}
end

DepthsFirst(x::T) where {T<:Node} = DepthsFirst{T}(Union{Child, T}[x])

function Base.iterate(x::DepthsFirst, state=x)
    isempty(x.remaining) && return nothing
    node = pop!(x.remaining)
    isleaf(node) || append!(x.remaining, node.children)
    return (node, x)
end

Base.IteratorEltype(::Type{<:DepthsFirst}) = Base.HasEltype()
Base.IteratorSize(::Type{<:DepthsFirst}) = Base.SizeUnknown()
Base.eltype(::Type{DepthsFirst{T}}) where T = Union{T, Child}

struct BreadthFirst{T}
    currentlevel::Vector{Union{Child, T}}
    nextlevel::Vector{Union{Child, T}}
end

BreadthFirst(x::T) where {T<:Node} = BreadthFirst{T}(Union{Child, T}[x], Union{Child, T}[])

function Base.iterate(x::BreadthFirst, state=x)
    if isempty(x.currentlevel)
        copy!(x.currentlevel, x.nextlevel)
        empty!(x.nextlevel)
    end
    isempty(x.currentlevel) && return nothing
    node = pop!(x.currentlevel)
    isleaf(node) || append!(x.nextlevel, node.children)
    return (node, x)
end

Base.IteratorEltype(::Type{<:BreadthFirst}) = Base.HasEltype()
Base.IteratorSize(::Type{<:BreadthFirst}) = Base.SizeUnknown()
Base.eltype(::Type{BreadthFirst{T}}) where T = Union{T, Child}

struct GenerationIterator{T}
    currentlevel::Vector{Union{Child, T}}
end

function GenerationIterator(x::T) where {T<:Node}
    return GenerationIterator{T}(Union{Child, T}[x])
end

function Base.iterate(x::GenerationIterator, state=x)
    isempty(x.currentlevel) && return nothing
    res = copy(x.currentlevel)
    empty!(x.currentlevel)
    for i in res
        append!(x.currentlevel, i.children)
    end
    return (res, x)
end

Base.IteratorEltype(::Type{<:GenerationIterator}) = Base.HasEltype()
Base.IteratorSize(::Type{<:GenerationIterator}) = Base.SizeUnknown()
Base.eltype(::Type{GenerationIterator{T}}) where T = Vector{Union{T, Child}}

generationsof(x::Node) = GenerationIterator(x)

nextgenerations(node::Node) = GenerationIterator{Child}(copy(node.children))

subtreeof(node::Node) = DepthsFirst(node)
descendantsof(node::Node) = DepthsFirst{Child}(copy(node.children))

isleaf(node::Node) = isempty(node.children)
leavesof(node::Child) = (child for child in subtreeof(node) if isleaf(child))
leavesof(node::Root) = (child for child in descendantsof(node) if isleaf(child))

function isredundant(node::Node)
    # Use descendantsof not subtree of because type stability
    length(node.children) == 1 && return true
    return any(length(i.children) == 1 for i in descendantsof(node))
end

function ispolytomic(node::Node)
    # Use descendantsof not subtree of because type stability
    length(node.children) > 2 && return true
    return any(length(i.children) > 2 for i in descendantsof(node))
end

function mrcaof(nodes::Vector{Root})
    isempty(nodes) && error("empty vector")
    root = first(nodes)
    length(nodes) == 1 && return root
    return any(nodes[i] !== root for i in 2:length(nodes)) ? nothing : root
end

function mrcaof(nodes::Vector{Node})
    length(nodes) == 1 && return first(nodes)
    roots = Vector{Root}()
    children = Vector{Children}()
    for node in nodes
        (node isa Root) ? push!(roots, node) : push!(children, node)
    end
    if length(roots) == 0
        return mrcaof(children)
    end
    root = first(roots)
    if length(roots) > 1 && any(roots[i] !== root for i in 2:length(roots))
        return nothing
    end
    if all(rootof(child) === root for child in children)
        return root
    else
        return nothing
    end
end

function mrcaof(nodes::Vector{Child})
    isempty(nodes) && error("empty vector")
    length(nodes) == 1 && return @inbounds nodes[1]
    length(nodes) == 2 && return mrcaof((@inbounds nodes[1]), (@inbounds nodes[2]))
    lineages = [collect(lineageof(node)) for node in nodes]
    minlength = minimum(map(length, lineages))
    generations = Array{Node}(undef, (length(nodes), minlength))

    for (i, lineage) in enumerate(lineages)
        generations[i,:] = lineage[end - minlength + 1 : end]
    end

    for generation in 1:minlength
        base = generations[1, generation]
        if all(i === base for i in generations[:, generation])
            return base
        end
    end
    return nothing
end

mrcaof(one::Node, two::Root) = rootof(one) === two ? two : nothing
mrcaof(one::Root, two::Child) = mrcaof(two, one)
function mrcaof(one::Child, two::Child)
    lin1, lin2 = collect(lineageof(one)), collect(lineageof(two))
    len1, len2 = length(lin1), length(lin2)
    smallen, largelen = minmax(len1, len2)
    if len1 != len2
        smallest, largest = ifelse(len1 < len2, (lin1, lin2), (lin2, lin1))
        unsafe_copyto!(largest, 1, largest, largelen-smallen+1, smallen)
    end
    @inbounds for i in 1:smallen
        if lin1[i] === lin2[i]
            return lin1[i]
        end
    end
    return nothing
end

treelength(node::Node) = sum(c.length for c in descendantsof(node))

function distance(one::Node, two::Node, mrca::Node)
    "Finds the distance between two nodes"
    dist = 0.0

    for leaf in (one, two)
        leaf isa Root && continue
        leafchild::Child = leaf
        for child in ChildLineageIterator(leafchild)
            child === mrca && break
            dist += child.length
        end
    end
    return dist
end

function distance(one::T1, two::T2) where {T1<:Node, T2<:Node}
    "Finds the distance between two nodes"
    mrca = mrcaof(one, two)
    if mrca === nothing
        error("the two nodes do not share an ancestor")
    end
    return distance(one, two, mrca)
end

function simplify!(node::T) where T<:Node
    "Deletes all redundant nodes in subtree, never deletes self"
    for child in descendantsof(node)
        length(child.children) == 1 && delete!(child)
    end
end

function subtree(nodes::Vector{<:Node})
    if isempty(nodes)
        error("no nodes given")
    elseif length(nodes) == 1
        return Root(first(nodes.name))
    end

    mrca = mrcaof(nodes)
    if mrca === nothing
        throw(ArgumentError("Not all nodes are related"))
    end
    root = Root(mrca.name)
    copyof = Dict{Node, Node}()
    copyof[mrca] = root

    # For the first node, create all the rest of the lineage
    for template in reverse!(collect(childlineageof(first(nodes))))
        copyof[template] = Child(template.name, copyof[template.parent], template.length)
    end

    # For each subsequent node
    lineage = Child[]
    for node in nodes[2:end]
        empty!(lineage)
        # While the ancestors have never been seen, keep getting more distanct
        # ancestors
        for template in childlineageof(node)
            if haskey(copyof, template)
                break
            end
            push!(lineage, template)
        end

        # Now instantiate all these unseen ancestors
        for template in reverse!(lineage)
            copyof[template] = Child(template.name, copyof[template.parent], template.length)
        end
    end

    simplify!(root)
    return root
end

# distance
# nodedistances

function _recursor(array, node, length)
    newlength = node.length + length
    push!(array, (node, newlength))
    for child in node.children
        _recursor(array, child, newlength)
    end
end

function _ndchild(node::Child, ancestors::Array{Child, 1}, withroot)
    pairs = Tuple{Node, Float64}[]

    dist = 0.0
    previousancestor = node
    for i in eachindex(ancestors)
        ancestor = ancestors[i]
        push!(pairs, (ancestor, dist))

        if i == 1
            for child in ancestor.children
                _recursor(pairs, child, dist)
            end
        else
            for child in filter(c->c != previousancestor, ancestor.children)
                _recursor(pairs, child, dist)
            end
        end

        dist += ancestor.length
        previousancestor = ancestor
    end

    # Add the root is asked to
    if withroot
        root::Root = ancestors[end].parent
        push!(pairs, (root, dist))

        for child in filter(c->c != ancestors[end], root.children)
            _recursor(pairs, child, dist)
        end
    end

    return pairs
end

nodedistances(node::Child) = _ndchild(node, collect(childlineageof(node)))

function nodedistances(node::Child, backto::Child)
    ancestors = collect(childlineageof(node))
    backtoindex = findfirst(isequal(backto), ancestors)
    if backtoindex === nothing
        error("backto must be an ancestor of node")
    end
    return _ndchild(node, ancestors[1:backtoindex], false)
end

function nodedistances(node::Child, backto::Root)
    ancestors = collect(childlineageof(node))
    if backto != ancestors[end].parent
        error("backto must be an ancestor of node")
    end
    return _ndchild(node, ancestors, true)
end

function nodedistances(node::Root)
    pairs = Tuple{Node, Float64}[(node, 0.0)]

    for child in node.children
        _recursor(pairs, child, 0)
    end

    return pairs
end

function nodedistances(node::Root, backto::Root)
    if backto !== node
        error("backto must be an ancestor of node")
    end
    return nodedistances(node)
end

function _furthestleaf(node, backto)
    leafdists::Array{Tuple{Child, Float64}} = filter(x->isleaf(x[1]), nodedistances(node, backto))
    maxdist = -1.0
    result = leafdists[1][1]
    for (leaf, dist) in leafdists
        if dist > maxdist
            result = leaf
        end
    end
    return result
end

function diameter(node::Node)
    leafone = _furthestleaf(node, node)
    leaftwo = _furthestleaf(leafone, node)

    return distance(leafone, leaftwo), leafone, leaftwo
end

@inline npairs(n::Int) = (n * (n-1)) >>> 1

function _update(node, cache)
    @inbounds begin
    nchildren = length(node.children)
    if nchildren == 0
        return (1, 0.0, 0.0)
    elseif nchildren == 1
        child = node.children[1]
        value = cache[child]
        return (value[1], value[2], value[3] + child.length)
    # These next two do the same thing, but the first one is faster for nchildren == 2
    elseif nchildren == 2
        return _update2(node, cache)
    else
        return _updaten(node, cache, nchildren)
    end
    end # inbounds
end

function _update2(node, cache)
    # See comments in _updaten
    @inbounds begin
    child1, child2 = node.children[1], node.children[2]
    L1, L2 = child1.length, child2.length
    v1, v2 = cache[child1], cache[child2]
    N1, N2 = v1[1], v2[1]
    I1, I2 = v1[2], v2[2]
    D1, D2 = v1[3], v2[3]

    N = N1 + N2
    I = (npairs(N1)*I1 + npairs(N2)*I2 + N1*N2*(D1+D2+L1+L2)) / (npairs(N1) + npairs(N2) + N1*N2)
    D = (N1*(D1+L1) + N2*(D2+L2)) / N
    end # inbounds
    return (N, I, D)
end

function _updaten(node, cache, nchildren)
    # We cache this to avoid looking up in the dictionary tonnes of times
    @inbounds begin
    childdata = Vector{Tuple{Int, Float64, Float64, Float64}}(undef, nchildren)
    for i in 1:nchildren
        child = node.children[i]
        value = cache[child]
        childdata[i] = (value[1], value[2], value[3], child.length)
    end

    # Here, we calculate the mean internal path distance to all the leaves
    sumd = 0.0
    npaths = 0
    for i in 1:nchildren
        # First add the paths within each child's leaves
        np = npairs(childdata[i][1])
        npaths += np
        sumd += np * childdata[i][2]
        # Then add the paths between the children
        for j in (i+1):nchildren
            np = childdata[i][1] * childdata[j][1]
            npaths += np
            sumd += np * (childdata[i][3] + childdata[i][4] + childdata[j][3] + childdata[j][4])
        end
    end
    internal = sumd / npaths

    # Calculate mean distances
    npaths = 0
    sumd = 0.0
    for i in 1:nchildren
        np = childdata[i][1]
        npaths += np
        sumd += np * (childdata[i][3] + childdata[i][4])
    end
    meandist = sumd / npaths
    end # inbounds
    return (npaths, internal, meandist)
end

function meandistance(node::Node)
    # We keep track of 3 properties for each node:
    # 1) Number of leaves descendants, 2) Mean internal distance between leaves
    # 3) Mean distance from leaf to the node
    cache = Dict{Child, Tuple{Int, Float64, Float64}}()
    generations = reverse!(collect(nextgenerations(node)))
    # We then iteratively merge siblings together WITHOUT needing to visit
    # all the sibling's descendants.
    for generation in generations
        for child in generation
            cache[child] = _update(child, cache)
        end
    end
    return _update(node, cache)[2]
end

function shufflenames!(node::Node)
    descendants = descendantsof(node)
    names = [child.name for child in descendants]
    shuffle!(names)

    @inbounds for i in eachindex(descendants)
        descendants[i].name = names[i]
    end

    return nothing
end

#= This function does two things:
1) Updated leafdistances for a node using the leafdistances of its children
2) Returns an array of one array per child, each containing the leaves of that child,
and the distance to the node. The distance between the leaves are then the sum of all leaves
across two of these inner arrays
=#
function _update_dist_parent!(node::Node, leafdistances::Dict{Node, Vector{Tuple{Int, Float64}}},
                              indexof::Dict{Child,Int})
    node_leafdistances = valtype(leafdistances)()
    dists_by_child = valtype(leafdistances)[]
    for child in node.children
        child_distances = valtype(leafdistances)()
        push!(dists_by_child, child_distances)

        # If child is leaf, it's not in leafdistances, but it leaves are just itself
        if isleaf(child)
            val = (indexof[child], child.length)
            push!(node_leafdistances, val)
            push!(child_distances, val)
        # Else, for each of its leaves, we add the length of the child
        else
            for (leafindex, dist) in leafdistances[child]
                val = (leafindex, dist + child.length)
                push!(node_leafdistances, val)
                push!(child_distances, val)
            end
            pop!(leafdistances, child)
        end
    end
    leafdistances[node] = node_leafdistances
    return dists_by_child
end

function _write_distances!(dists::Matrix{Float64}, dists_by_child::Vector{Vector{Tuple{Int, Float64}}})
    @inbounds for (vec1, vec2) in Pairs(dists_by_child)
        for (index1, dist1) in vec1, (index2, dist2) in vec2
            # This is distance from leaf A to node + distance from leaf B to node
            dist = dist1 + dist2
            dists[index1, index2] = dist
            dists[index2, index1] = dist
        end
    end
end

#= Algorithm: Start from youngest leaves, moving up to the MRCA. Along the way,
keep track of the distance to the node you're at from all its leaf descendants.
Every time you reach a new node, that node is the MRCA of leaf descendants of two different
children, and you already got the distance from those leaves to the node, so just add them.
=#
function distance_matrix(nodes::Vector{Child})
    all(isleaf, nodes) || throw(ArgumentError("All nodes must be leaves"))
    # This dict is node => [(index, dist_from_leaf_to_node) ..] for all node's leaves
    leafdistances = Dict{Node, Vector{Tuple{Int, Float64}}}()
    mrca = mrcaof(nodes)
    generations = reverse!(collect(generationsof(mrca)))
    result = zeros(length(nodes), length(nodes))
    indexof = Dict(i=>n for (n,i) in enumerate(nodes))
    # Start at youngest generation, skipping first that only has leaves
    for generation in generations[2:end]
        for node in generation
            if !isleaf(node)
                dists_by_child = _update_dist_parent!(node, leafdistances, indexof)
                _write_distances!(result, dists_by_child)
            end
        end
    end
    return result
end

reroot!(root::Root) = root
function reroot!(node::Child)
    lineage = collect(childlineageof(node))
    lengths = [node.length for node in lineage]
    root::Root = lineage[end].parent

    # Switch properties of node and root
    root.name, node.name = node.name, root.name
    oldrootchildren = copy(root.children)
    root.children = node.children
    push!(root.children, node.parent)
    for child in root.children
        child.parent = root
    end

    node.parent = lineage[end]
    node.length = lengths[end]
    deleteat!(oldrootchildren, findfirst(isequal(lineage[end]), oldrootchildren))
    node.children = oldrootchildren

    for child in node.children
        child.parent = node
    end

    for childindex in 2:length(lineage)
        child = lineage[childindex]
        # Set length
        child.length = lengths[childindex - 1]

        # Set new parent
        if childindex == 2
            child.parent = root
        else
            child.parent = lineage[childindex - 1]
        end

        # Identify new child
        if childindex == length(lineage)
            newchild = node
        else
            newchild = lineage[childindex + 1]
        end

        # Identify old child to be removed and replace old with new
        oldchild = lineage[childindex - 1]
        @assert oldchild in child.children
        i = findfirst(isequal(oldchild), child.children)
        child.children[i] = newchild
    end

    return root
end

function reroot!(node::Child, length::Number)
    root = rootof(node)
    newnode = insert(root.name, node, length)
    newroot = reroot!(newnode)

    if length(newnode.children) == 1
        delete!(newnode)
    end

    return newroot
end

function midpointof(node::Node)
    "Finds the midpoint - i.e. halfway through the diameter"

    diameter, leafone, leaftwo = diameter(node)
    radius = diameter / 2
    mrca = mrcaof(leafone, leaftwo)

    if distance(leafone, mrca, mrca) >= radius
        longestleaf = leafone
    else
        longestleaf = leaftwo
    end

    traveled = 0.0
    for child in childlineageof(longestleaf)
        traveled += child.length

        if traveled > radius
            return child, radius + child.length - traveled
        end
    end

    error("Should not reach this line")
end

function reroot_midpoint!(node::Root)
    "Reroots the tree at the midpoint."

    child, distance = midpointof(node)
    newnode = insert(node.name, child, distance)
    newroot = reroot!(newnode)
    delete!(newnode)
    return newroot
end

struct CompactTree
    rootname::String
    names::Vector{String}
    parentids::Vector{Int}
    lengths::Vector{Float64}
end

function pack(node::Node)
    nodemap = Dict{Node,Int}(node => 0)
    names = Vector{String}()
    parents = Vector{Int}()
    lengths = Vector{Float64}()
    for (n, child) in enumerate(descendantsof(node))
        push!(names, child.name)
        push!(lengths, child.length)
        parentid = nodemap[child.parent]
        push!(parents, parentid)
        nodemap[child] = n
    end
    return CompactTree(node.name, names, parents, lengths)
end

function unpack(c::CompactTree)
    root = Root(c.rootname)
    nodemap = Dict{Int, Node}(0 => root)
    for i in eachindex(c.names)
        parent = nodemap[c.parentids[i]]
        child = Child(c.names[i], parent, c.lengths[i])
        nodemap[i] = child
    end
    reverse!(root.children)
    for child in descendantsof(root)
        isleaf(child) || reverse!(child.children)
    end
    return root
end

include("parsetree.jl")

export Node, Child, Root,
    # Instantiations
    newick,  json, deepcopyasroot, pack, unpack,
    # Single-node mutations
    delete!, insert, detach!, move!,
    # Transversal
    lineageof, childlineageof, rootof, descendantsof, subtreeof, mrcaof,
    DepthsFirst, BreadthFirst, generationsof, nextgenerations, leavesof,
    # Boolean
    isredundant, ispolytomic, isleaf,
    # Metrics
    distance, nodedistances, diameter, midpointof, treelength, meandistance,
    distance_matrix,
    # Whole-tree rearrangements
    simplify!, shufflenames!, reroot!, reroot_midpoint!,
    # Methods returning new trees
    subtree,
    # Misc
    namemapof

end # module Tree

max_stars_repo_name: jakobnissen/tree

max_issues_repo_name: jakobnissen/tree

max_forks_repo_name: jakobnissen/tree


text:
# Model
include_model("hopper")
model_ft = free_time_model(model)

# Horizon
T = 51

# Time step
tf = 1.0
h = tf / (T - 1)

# Bounds
_uu = Inf * ones(model_ft.m)
_uu[model_ft.idx_u] = model_ft.uU
_uu[end] = 2.0 * h
_ul = zeros(model_ft.m)
_ul[model_ft.idx_u] = model_ft.uL
_ul[end] = 0.5 * h
ul, uu = control_bounds(model_ft, T, _ul, _uu)

# Initial and final states
z_h = 0.25
q1 = [0.0, 0.5 + z_h, 0.0, 0.25]

xl, xu = state_bounds(model_ft, T,
		[model_ft.qL; model_ft.qL],
		[model_ft.qU; model_ft.qU],
        x1 = [q1; Inf * ones(model.nq)])

# Objective
include_objective("velocity")
obj_tracking = quadratic_time_tracking_objective(
    [Diagonal(ones(model_ft.n)) for t = 1:T],
    [Diagonal([1.0, 1.0e-1,
		zeros(model_ft.nc)..., zeros(model_ft.nb)...,
		zeros(model_ft.m - model_ft.nu - model_ft.nc - model_ft.nb - 1)..., 0.0])
		for t = 1:T-1],
    [[q1; q1] for t = 1:T],
    [zeros(model_ft.m) for t = 1:T-1],
    1.0)

include_objective("control_velocity")
obj_contact_penalty = PenaltyObjective(1.0e5, model_ft.m - 1)
obj_velocity = velocity_objective(
    [1.0e-1 * Diagonal([1.0, 1.0, 1.0, 1.0]) for t = 1:T-1],
    model_ft.nq,
    h = 1.0,
    idx_angle = collect([3]))
obj_control_velocity = control_velocity_objective(Diagonal([1.0; 1.0; zeros(model_ft.m - 2)]))
obj = MultiObjective([obj_tracking,
	obj_contact_penalty,
	obj_velocity,
	obj_control_velocity])

# Constraints
include_constraints(["free_time", "contact", "loop"])
con_free_time = free_time_constraints(T)
con_contact = contact_constraints(model_ft, T)
con_loop = loop_constraints(model, (1:model.n), 1, T)
con = multiple_constraints([con_free_time, con_contact, con_loop])

# Problem
prob = trajectory_optimization_problem(model_ft,
               obj,
               T,
               xl = xl,
               xu = xu,
               ul = ul,
               uu = uu,
               con = con)

# Trajectory initialization
x0 = [[q1; q1] for t = 1:T] # linear interpolation on state
u0 = [[1.0e-5 * rand(model_ft.m-1); h] for t = 1:T-1] # random controls

# Pack trajectories into vector
z0 = pack(x0, u0, prob)

#NOTE: may need to run examples multiple times to get good trajectories
# Solve nominal problem

include_snopt()
@time z̄, info = solve(prob, copy(z0),
	nlp = :SNOPT7,
	tol = 1.0e-5, c_tol = 1.0e-5, mapl = 5,
	time_limit = 60)
@show check_slack(z̄, prob)
x̄, ū = unpack(z̄, prob)
tf, t, h̄ = get_time(ū)

z_traj = [x[2] for x in x̄]
λ_traj = [u[model_ft.idx_λ] for u in ū]
f_traj = [u[2] for u in ū]

using Plots
plot(h̄[1] * hcat(λ_traj...)', linetype = :steppost, label = ["" "λ1"])
plot!(h̄[1] * hcat(f_traj...)', linetype = :steppost, label = "force")

plot(hcat(z_traj...)', linetype = :steppost, label = "z")
plot(hcat(ū...)[1:model.nu, :]', linetype = :steppost, label = "u")

include(joinpath(pwd(), "models/visualize.jl"))
vis = Visualizer()
render(vis)
visualize!(vis, model_ft,
	# state_to_configuration(x̄),
	state_to_configuration([x̄..., ([x̄[2:end] for i = 1:3]...)...]),
	Δt = h̄[1],
	scenario = :vertical)

# Model-predictive control
include_snopt()
# include(joinpath(pwd(), "src/contact_simulator/simulator_variable_time_step.jl"))
include(joinpath(pwd(), "src/contact_simulator/simulator.jl"))

function run_mpc()
	# Time steps
	h_mpc = h̄[1] # mpc time step
	d_sim = 10 # simulator rate
	h_sim = h_mpc / d_sim # simulator time step

	# MPC horizon
	T_mpc =  3 * T # mpc horizon

	# Reference trajectories
	n_gaits = 20
	T_gait = n_gaits * T - (n_gaits - 1)
	_ū = [u[1:end-1] for u in ū] # remove time step control input
	x_ref = [x̄..., ([x̄[2:end] for i = 1:n_gaits-1]...)...]
	u_ref = [_ū..., ([_ū for i = 1:n_gaits-1]...)...]
	u_ref_ft = [ū..., ([ū for i = 1:n_gaits-1]...)...]

	Q_mpc = Diagonal([10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0])
	R_mpc = Diagonal([1.0e-1 * ones(model.nu); 1.0e-1 * ones(model.m - model.nu)])
	R_mpc_ft = Diagonal([1.0e-1 * ones(model.nu); 1.0e-1 * ones(model.m - model.nu); 0.0])

	_, P_mpc = tvlqr(model_ft, x_ref, u_ref_ft, h_mpc,
		[Q_mpc for t = 1:T_gait], [R_mpc_ft for t = 1:T_gait-1])

	x1 = x_ref[1] + 1.0e-8 * randn(model.n)
    x_hist = [x1]
	q_sim = [x1[1:model.nq], x1[model.nq .+ (1:model.nq)]]
	v_sim = [(q_sim[end] - q_sim[end-1]) ./ h_mpc]
    u_hist = []

	T_sim = 400
	shift = 1

    for t = 1:T_sim-1

		# create MPC problem
    	println("T = $t")
		_uu = Inf * ones(model.m)
		_uu[model_ft.idx_u] = model.uU
		_ul = zeros(model.m)
		_ul[model.idx_u] = model.uL
		ul_mpc, uu_mpc = control_bounds(model, T_mpc, _ul, _uu)

		xl_mpc, xu_mpc = state_bounds(model, T_mpc,
				[model.qL; model.qL],
				[model.qU; model.qU],
		        x1 = x_hist[end])

    	# Objective
    	obj_mpc = quadratic_tracking_objective(
    	        [t < T_mpc ? Q_mpc : P_mpc[T_mpc + shift - 1] for t = 1:T_mpc],
				# [Q_mpc for t = 1:T_mpc],
    	        [R_mpc for t = 1:T_mpc-1],
    	        [x_ref[shift + t - 1] for t = 1:T_mpc],
    			[u_ref[shift + t - 1] for t = 1:T_mpc-1])

    	# Problem
    	prob_mpc = trajectory_optimization_problem(model,
    			obj_mpc,
    			T_mpc,
    			h = h_mpc,
    			ul = ul_mpc,
    			uu = uu_mpc,
    			xl = xl_mpc,
    			xu = xu_mpc)

    	# Pack trajectories into vector
    	z0_mpc = pack(x_ref[shift .+ (1:T_mpc)], u_ref[shift .+ (1:T_mpc-1)], prob_mpc)

    	# Solve MPC problem
    	@time z_mpc, info = solve(prob_mpc, copy(z0_mpc),
            nlp = :SNOPT7, mapl = 0, tol = 1.0e-3, c_tol = 1.0e-3)
    	x_mpc, u_mpc = unpack(z_mpc, prob_mpc)

		u_input = max.(min.(u_mpc[1], uu_mpc[1]), ul_mpc[1])[1:model.nu] # return feasible controls
		# u_input = u_mpc[1][1:model.nu]
		# simulate
		f_dist = 100.0
		w0 = t == T ? [f_dist, 0.0, 0.0, 0.0] : 0.0e-8 * randn(model.d) # disturbance
		# u_input = u_ref[shift][1:model.nu]
		# w0 = zeros(model.d)
		# push!(x_hist, step_contact(model, x_hist[end],
		# 		u_input, w0, h_mpc,
		# 		tol_c = 1.0e-5, tol_opt = 1.0e-5, tol_s = 1.0e-4,
		# 		nlp = :SNOPT7))

		for i = 1:d_sim
			# step simulator
			q = step_contact(model,
			   v_sim[end], q_sim[end-1], q_sim[end],
			   u_input, w0, h_sim)
			v = (q - q_sim[end]) / h_sim

			push!(q_sim, q)
			push!(v_sim, v)
		end

		push!(x_hist, [x_hist[end][model.nq .+ (1:model.nq)]; q_sim[end]])
    	push!(u_hist, u_input)
    	shift += 1
    end

    return x_hist, u_hist, x_ref[1:T_sim], u_ref[1:T_sim-1]
end

@time x_hist_mpc, u_hist_mpc, x_ref, u_ref = run_mpc()

plot(hcat(state_to_configuration(x_ref)...)',
	label = "", color = :black, width = 2.0)
plot!(hcat(state_to_configuration(x_hist_mpc)...)',
	label = "", color = :magenta, width = 1.0)

plot(hcat(u_ref...)[1:model.nu,:]',
	label = "u (ref)", color = :black, width = 2.0)
plot!(hcat(u_hist_mpc...)[1:model.nu,:]',
	label = "u (mpc)", color = :magenta, width = 1.0)

include(joinpath(pwd(), "models/visualize.jl"))
vis = Visualizer()
render(vis)
visualize!(vis, model, x_hist_mpc, Δt = h̄[1], scenario = :push)
open(vis)


text:
# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md

typealias Node Vector{Float64}

abstract AbstractPoint

type Point{P<:AbstractPoint}
    id :: Int
    weight :: Float64
    coords :: Vector{Float64}
    fields :: Dict{AbstractString, Field}
    properties :: P
end

function setindex!{T}(point::Point, val::Pair{Float64, T}, field_name)
    point.fields[field_name] = Field(val)
end

function getindex(point::Point, field_name)
    return point.fields[field_name]
end

function getindex(point::Point, idx::Int)
    return point.coords[idx]
end

function haskey(point::Point, field_name)
    return haskey(point.fields, field_name)
end

function call(point::Point, field_name, time=0.0)
    point.fields[field_name](time).data
end

function update!{T}(point::Point, field_name, val::Pair{Float64, T})
    if haskey(point, field_name)
        update!(point[field_name], val)
    else
        point[field_name] = val
    end
end

#= TODO: in future
type Node <: AbstractPoint
end

type MaterialPoint <: AbstractPoint
end
=#

type IntegrationPoint <: AbstractPoint
end

typealias IP Point{IntegrationPoint}

function IP(id, weight, coords)
    return IP(id, weight, coords, Dict(), IntegrationPoint())
end

function convert(::Type{IP}, data::Tuple{Float64, Vector{Float64}})
    weight, coords = data
    return IP(-1, weight, coords)
end


text:
#=
Popis
  Funkcia vytvori histogram zaznamenavajuci pocetnost stupna vrcholov v grafe.

Parametre
  g : datova struktura grafu :: GenericGraph

Navratova hodnota
  histogram : slovnik s mapovanim kluc=stupen_vrchola => hodnota=pocetnost_vrcholov :: Dict{Int64, Int64}()

Poznamky
=#

function degree_histogram(g::GenericGraph)

  histogram = [out_degree(v, g) => 0 for v = vertices(g)]

  for v in vertices(g)
    histogram[out_degree(v, g)] += 1
  end

  return histogram

end

max_stars_repo_name: mpuk/GraphAnalysis

max_issues_repo_name: mpuk/GraphAnalysis

max_forks_repo_name: mpuk/GraphAnalysis


text:
# AUTO GENERATED FILE - DO NOT EDIT

export 'feffery'_antdtitle

"""
    'feffery'_antdtitle(;kwargs...)
    'feffery'_antdtitle(children::Any;kwargs...)
    'feffery'_antdtitle(children_maker::Function;kwargs...)

An AntdTitle component.

Keyword arguments:
- `children` (a list of or a singular dash component, string or number; optional): The content of the tab - will only be displayed if this tab is selected
- `id` (String; optional)
- `className` (String; optional)
- `code` (Bool; optional)
- `copyable` (Bool; optional)
- `disabled` (Bool; optional)
- `italic` (Bool; optional)
- `keyboard` (Bool; optional)
- `level` (Real; optional)
- `loading_state` (optional): . loading_state has the following type: lists containing elements 'is_loading', 'prop_name', 'component_name'.
Those elements have the following types:
  - `is_loading` (Bool; optional): Determines if the component is loading or not
  - `prop_name` (String; optional): Holds which property is loading
  - `component_name` (String; optional): Holds the name of the component that is loading
- `locale` (a value equal to: 'zh-cn', 'en-us'; optional)
- `mark` (Bool; optional)
- `strikethrough` (Bool; optional)
- `strong` (Bool; optional)
- `style` (Dict; optional)
- `type` (a value equal to: 'secondary', 'success', 'warning', 'danger'; optional)
- `underline` (Bool; optional)
"""
function 'feffery'_antdtitle(; kwargs...)
        available_props = Symbol[:children, :id, :className, :code, :copyable, :disabled, :italic, :keyboard, :level, :loading_state, :locale, :mark, :strikethrough, :strong, :style, :type, :underline]
        wild_props = Symbol[]
        return Component("'feffery'_antdtitle", "AntdTitle", "feffery_antd_components", available_props, wild_props; kwargs...)
end

'feffery'_antdtitle(children::Any; kwargs...) = 'feffery'_antdtitle(;kwargs..., children = children)
'feffery'_antdtitle(children_maker::Function; kwargs...) = 'feffery'_antdtitle(children_maker(); kwargs...)

max_stars_repo_name: ReilYoung/feffery-antd-components

max_issues_repo_name: ReilYoung/feffery-antd-components

max_forks_repo_name: ReilYoung/feffery-antd-components


text:
#=##############################################################################
# DESCRIPTION
    Tools for post-processing outputs from UCD-QuietFly

# AUTHORSHIP
  * Author    : Eduardo J. Alvarez
  * Email     : Edo.AlvarezR@gmail.com
  * Created   : Apr 2021
  * License   : MIT
=###############################################################################

# """
# Reads QuietFly's output files under `read_path` and formats them into TecPlot
# format, saving them under `save_path`.
# """
# function quietfly2tec(read_path, save_path)
#
# end

function quietfly_plot_spectrum_spl(read_path, run_name, mici; onethirdoctave=true)

end

max_stars_repo_name: byuflowlab/FLOWNoise

max_issues_repo_name: byuflowlab/FLOWNoise

max_forks_repo_name: byuflowlab/FLOWNoise


text:
# --------------------------------------------------------------------------
# ACE.jl and SHIPs.jl: Julia implementation of the Atomic Cluster Expansion
# Copyright (c) 2019 Christoph Ortner <christophortner0@gmail.com>
# Licensed under ASL - see ASL.md for terms and conditions.
# --------------------------------------------------------------------------

function evaluate!(A, tmp, basis::OneParticleBasis, Rs, Zs::AbstractVector, z0)
   fill!(A, 0)
   iz0 = z2i(basis, z0)
   for (R, Z) in zip(Rs, Zs)
      iz = z2i(basis, Z)
      add_into_A!((@view A[basis.Aindices[iz, iz0]]), tmp, basis, R, iz, iz0)
   end
   return A
end

function evaluate_d!(A, dA, tmpd, basis::OneParticleBasis,
                     Rs, Zs::AbstractVector, z0)
   fill!(A, 0)
   fill!(dA, zero(eltype(dA)))  # TODO: this should not be necessary!
   iz0 = z2i(basis, z0)
   for (j, (R, Z)) in enumerate(zip(Rs, Zs))
      iz = z2i(basis, Z)
      Aview = @view A[basis.Aindices[iz, iz0]]
      dAview = @view dA[basis.Aindices[iz, iz0], j]
      add_into_A_dA!(Aview, dAview, tmpd, basis, R, iz, iz0)
   end
   return dA
end

function evaluate!(A, tmp, basis::OneParticleBasis, R, z::AtomicNumber, z0)
   fill!(A, 0)
   iz0, iz = z2i(basis, z0), z2i(basis, z)
   add_into_A!((@view A[basis.Aindices[iz, iz0]]), tmp, basis, R, iz, iz0)
   return A
end

function evaluate_d!(A, dA, tmpd, basis::OneParticleBasis,
                     R, z::AtomicNumber, z0)
   fill!(A, 0)
   iz, iz0 = z2i(basis, z), z2i(basis, z0)
   Aview = @view A[basis.Aindices[iz, iz0]]
   dAview = @view dA[basis.Aindices[iz, iz0]]
   add_into_A_dA!(Aview, dAview, tmpd, basis, R, iz, iz0)
   return dA
end

function alloc_B(basis::OneParticleBasis, args...)
   NZ = numz(basis)
   maxlen = maximum( sum( length(basis, iz, iz0) for iz = 1:NZ )
                     for iz0 = 1:NZ )
   T = fltype(basis)
   return zeros(T, maxlen)
end

# TODO: fix the dispatch structure on alloc_dB
alloc_dB(basis::OneParticleBasis, ::JVec) = alloc_dB(basis, 1)
alloc_dB(basis::OneParticleBasis, Rs::AbstractVector{<: JVec}, args...) =
      alloc_dB(basis, length(Rs))
alloc_dB(basis::OneParticleBasis) = alloc_dB(basis, 1)

function alloc_dB(basis::OneParticleBasis, maxN::Integer)
   NZ = numz(basis)
   maxlen = maximum( sum( length(basis, iz, iz0) for iz = 1:NZ )
                     for iz0 = 1:NZ )
   T = fltype(basis)
   return zeros(JVec{T}, (maxlen, maxN))
end

function set_Aindices!(basis::OneParticleBasis)
   NZ = numz(basis)
   for iz0 = 1:NZ
      idx = 0
      for iz = 1:NZ
         len = length(basis, iz, iz0)
         basis.Aindices[iz, iz0] = (idx+1):(idx+len)
         idx += len
      end
   end
   return basis
end

get_basis_spec(basis::OneParticleBasis, iz0::Integer) =
      get_basis_spec(basis, i2z(iz0))

get_basis_spec(basis::OneParticleBasis, s::Symbol) =
      get_basis_spec(basis, atomic_number(s))


text:
using NgSpice
using Test

NgSpice.init()

nl = """
* Diode test circuit
V1 inp 0 DC=1 AC=1 SIN(1 1 1000)
* Connected in forward bias:
D1 inp 0 dmod
.model dmod d cjo=0 m=1 rs=5000 area=1 pj=0 cjp=0 mj=0 mjsw=0 pb=1 php=1 n=2 bv=0 ibv=0 ik=0 ikr=0 is=1.0e-9 tt=0
.save all @d1[id] @d1[cd] @d1[charge] @d1[vd]
.END
"""

@test NgSpice.load_netlist(nl) == 0

max_stars_repo_name: JuliaComputing/Ngspice

max_issues_repo_name: JuliaComputing/Ngspice

max_forks_repo_name: JuliaComputing/Ngspice