text:
# using Continuables
# import ResumableFunctions: @resumable, @yield
using ResumableFunctions

"""
    evolutionOperator(Hamiltonian, t)

Get evolution operator as Operator using the definition

``U(t) = e^{-i H t / \\hbar}, \\quad \\hbar = 1``.

"""
function evolutionOperator(Hamiltonian::Operator, t::AbstractFloat)::Operator
    return exp(-1im * Hamiltonian * t)
end

"""
    evolutionOperatorA(H_lambda, S, Sinv, t)

Get evolution operator as Array using the definition

``U(t) = S e^{-i H_\\lambda t / \\hbar} S^{-1},``

where ``\\hbar = 1``, ``H_{\\lambda,ii} = \\lambda_i`` are eigenvalues of ``H``,
so ``H`` has to be non-singular, otherwise ``H_{\\lambda,ij} = 0, i \\neq j``.
``S`` is obtained from eigendecomposition of ``H``, for example

```julia
H_lambda, S = eigen(Hamiltonian.data)
Sinv = inv(S)
H_lambda = diagm(H_lambda)
```
and arguments have to be Arrays.

"""
function evolutionOperatorA(H_lambda::Array, S::Array, Sinv::Array, t::AbstractFloat)::Array
    return S * exp(-1im * H_lambda * t) * Sinv
end

"""
    evolutionSuperOperator(Hamiltonian, t)

Get evolution operator as SuperOperator using the definition

``\\mathcal{U}(t) \\:\\cdot\\: = U(t) \\:\\cdot\\: U^\\dagger(t) = e^{-i H t / \\hbar} \\:\\cdot\\: e^{i H t / \\hbar}, \\quad \\hbar = 1``.

"""
function evolutionSuperOperator(Hamiltonian::Operator, t::AbstractFloat)::SuperOperator
    U = evolutionOperator(Hamiltonian, t)
    return spre(U) * spost(U')
end

"""
    evolutionOperatorArray(Hamiltonian, tspan)

Get evolution operators as Vector or Operators using the definition

``U(t) = e^{-i H t / \\hbar}, \\quad \\hbar = 1``.

"""
function evolutionOperatorArray(Hamiltonian::Operator, tspan::Array)::Array
    N = length(tspan)
    Ham_lambda, S = eigen(Hamiltonian.data)
    Sinv = inv(S)
    U_diagonal = zeros(ComplexF64, size(Ham_lambda))

    U_ref = evolutionOperator(Hamiltonian, 0.0)
    U_op_array = [deepcopy(U_ref)]
    for t_i = 2:N
        push!(U_op_array, deepcopy(U_ref))
    end

    for t_i = 1:N
        U_diagonal .= map(lambda -> exp(-1im * lambda * tspan[t_i]), Ham_lambda)
        U = S * diagm(U_diagonal) * inv(S)
        U_op_array[t_i].data[:, :] .= U
    end
    return U_op_array
end

"""
    evolutionOperatorArray(Hamiltonian, tspan)

Get evolution superoperators as Vector or SuperOperators using the definition

``\\mathcal{U}(t) \\:\\cdot\\: = U(t) \\:\\cdot\\: U^\\dagger(t) = e^{-i H t / \\hbar} \\:\\cdot\\: e^{i H t / \\hbar}, \\quad \\hbar = 1``.

"""
function evolutionSuperOperatorArray(Hamiltonian::Operator, tspan::Array)::Array
    N = length(tspan)
    Ham_lambda, S = eigen(Hamiltonian.data)
    basis = GenericBasis([size(Ham_lambda, 1)])
    Sinv = inv(S)
    U_diagonal = zeros(ComplexF64, size(Ham_lambda))

    UU_ref = evolutionSuperOperator(Hamiltonian, 0.0)
    U_supop_array = [deepcopy(UU_ref)]
    for t_i = 2:N
        push!(U_supop_array, deepcopy(UU_ref))
    end

    for t_i = 1:N
        U_diagonal .= map(lambda -> exp(-1im * lambda * tspan[t_i]), Ham_lambda)
        U = S * diagm(U_diagonal) * inv(S)
        U_op = DenseOperator(basis, basis, U)
        U_supop_array[t_i] = spre(U_op) * spost(U_op')
    end
    return U_supop_array
end

"""
    evolutionOperatorIterator(Hamiltonian, tspan; diagonalize = true, approximate = false)

Resumable function that returns evolution operator as Operator type at the time t from tspan.
See [`evolutionOperator`](@ref). The `diagonalize` argument decompose Hamiltonian into ``\\lambda_i, S, S^{-1}``
and calculate the exponential using eigenvalue decomposition. The `approximate` option asusmes that tspan is made of
equidistant points, therefore ``U(t)`` has to be calculated for ``t_0`` and ``t_\\text{step}``.

"""
function evolutionOperatorIterator end

@resumable function evolutionOperatorIterator(
    Hamiltonian::Operator,
    tspan::Array;
    diagonalize = true,
    approximate = false,
)
    N = length(tspan)
    t_step = tspan[2] - tspan[1]
    if diagonalize
        Ham_lambda, S = eigen(Hamiltonian.data)
        basis = GenericBasis([size(Ham_lambda, 1)])
        Sinv = inv(S)

        U_diagonal = zero(Ham_lambda)
        U_diagonal = map(lambda -> exp(-1im * lambda * tspan[1]), Ham_lambda)
        U = S * diagm(U_diagonal) * inv(S)
        U_op = DenseOperator(basis, basis, U)
    elseif !approximate
        U_op = evolutionOperator(Hamiltonian, tspan[1])
    else
        U_op = evolutionOperator(Hamiltonian, tspan[1])
        U_step = evolutionOperator(Hamiltonian, t_step)
    end
    @yield U_op

    for t_i = 2:N
        t = tspan[t_i]
        if diagonalize
            U_diagonal .= map(lambda -> exp(-1im * lambda * t), Ham_lambda)
            U_op.data .= S * diagm(U_diagonal) * inv(S)
        elseif !approximate
            U_op = evolutionOperator(Hamiltonian, t)
        else
            U_op = U_step * U_op
        end
        @yield U_op
    end
end

"""
    evolutionSuperOperatorIterator(Hamiltonian, tspan;
    \tdiagonalize = true, approximate = false)

Resumable function that returns evolution operator as Operator type at the time t from tspan.
See [`evolutionSuperOperator`](@ref). The `diagonalize` argument decompose Hamiltonian into ``\\lambda_i, S, S^{-1}``
and calculate the exponential using eigenvalue decomposition. The `approximate` option asusmes that tspan is made of
equidistant points, therefore ``U(t)`` has to be calculated for ``t_0`` and ``t_\\text{step}``.

"""
function evolutionSuperOperatorIterator end

@resumable function evolutionSuperOperatorIterator(
    Hamiltonian::Operator,
    tspan::Array;
    diagonalize = true,
    approximate = false,
)
    N = length(tspan)
    t_step = tspan[2] - tspan[1]
    if diagonalize
        Ham_lambda, S = eigen(Hamiltonian.data)
        basis = GenericBasis([size(Ham_lambda, 1)])
        Sinv = inv(S)

        U_diagonal = zero(Ham_lambda)
        U_diagonal = map(lambda -> exp(-1im * lambda * tspan[1]), Ham_lambda)
        U = S * diagm(U_diagonal) * Sinv
        U_op = DenseOperator(basis, basis, U)
    elseif !approximate
        U_op = evolutionOperator(Hamiltonian, tspan[1])
    else
        U_op = evolutionOperator(Hamiltonian, tspan[1])
        U_step = evolutionOperator(Hamiltonian, t_step)
    end
    U_supop = spre(U_op) * spost(U_op')
    @yield U_supop

    for t_i = 2:N
        t = tspan[t_i]
        if diagonalize
            U_diagonal .= map(lambda -> exp(-1im * lambda * tspan[t_i]), Ham_lambda)
            U_op.data .= S * diagm(U_diagonal) * Sinv
            U_supop.data .= (spre(U_op)).data * (spost(U_op')).data
        elseif !approximate
            U_op = evolutionOperator(Hamiltonian, t)
            U_supop = spre(U_op) * spost(U_op')
        else
            U_op = U_step * U_op
            U_supop = spre(U_op) * spost(U_op')
        end
        @yield U_supop
    end
end

"""
    evolutionExact(ket0, tspan, Hamiltonian; diagonalize = true, approximate = false)

Calculate exact time evolution of the `ket0` state see [`evolutionOperatorIterator`](@ref).
The `diagonalize` argument decompose Hamiltonian into ``\\lambda_i, S, S^{-1}``
and calculate the exponential using eigenvalue decomposition. The `approximate` option asusmes that tspan is made of
equidistant points, therefore ``U(t)`` has to be calculated for ``t_0`` and ``t_\\text{step}``.

"""
function evolutionExact(
    ket0::Ket,
    tspan::Array,
    Hamiltonian::Operator;
    diagonalize = true,
    approximate = false,
)
    N = length(tspan)
    ket_array = Array{Ket,1}(undef, 0)
    for t_i = 1:N
        push!(ket_array, deepcopy(ket0))
    end
    t_i = 0
    for U_op in evolutionOperatorIterator(
        Hamiltonian,
        tspan;
        diagonalize = diagonalize,
        approximate = approximate,
    )
        t_i += 1
        ket_array[t_i] = U_op * ket0
    end
    return ket_array
end

"""
    evolutionExact!(ket_array, ket0, tspan, Hamiltonian;
    \tdiagonalize = true, approximate = false)

Calculate exact time evolution of the `ket0` state inplace see [`evolutionOperatorIterator`](@ref).
The `diagonalize` argument decompose Hamiltonian into ``\\lambda_i, S, S^{-1}``
and calculate the exponential using eigenvalue decomposition. The `approximate` option asusmes that tspan is made of
equidistant points, therefore ``U(t)`` has to be calculated for ``t_0`` and ``t_\\text{step}``.

"""
function evolutionExact!(
    ket_array::Array{Array{C,1},1},
    ket0::Ket,
    tspan::Array,
    Hamiltonian::Operator;
    diagonalize = true,
    approximate = false,
) where {C<:ComputableType}
    N = length(tspan)
    buffer = zeros(C, size(ket0.data))
    for t_i = 1:N
        push!(ket_array, deepcopy(ket0.data))
    end
    t_i = 0
    for U_op in evolutionOperatorIterator(
        Hamiltonian,
        tspan;
        diagonalize = diagonalize,
        approximate = approximate,
    )
        t_i += 1
        buffer[:] .= convert(Array{C,1}, (U_op * ket0).data)
        ket_array[t_i] = deepcopy(buffer)
    end
    return ket_array
end

"""
    evolutionExact(op0, tspan, Hamiltonian; diagonalize = true, approximate = false)

Calculate exact time evolution of the `op0` state see [`evolutionOperatorIterator`](@ref).
The `diagonalize` argument decompose Hamiltonian into ``\\lambda_i, S, S^{-1}``
and calculate the exponential using eigenvalue decomposition. The `approximate` option asusmes that tspan is made of
equidistant points, therefore ``U(t)`` has to be calculated for ``t_0`` and ``t_\\text{step}``.

"""
function evolutionExact(
    op0::Operator,
    tspan::Array,
    Hamiltonian::Operator;
    diagonalize = true,
    approximate = false,
)
    N = length(tspan)
    op_array = Array{typeof(op0),1}(undef, 0)
    for t_i = 1:N
        push!(op_array, deepcopy(op0))
    end
    t_i = 0
    for U_op in evolutionOperatorIterator(
        Hamiltonian,
        tspan;
        diagonalize = diagonalize,
        approximate = approximate,
    )
        t_i += 1
        op_array[t_i] = U_op * op0 * U_op'
    end
    return op_array
end

"""
    evolutionExact!(op_array, op0, tspan, Hamiltonian;
    \tdiagonalize = true, approximate = false)

Calculate exact time evolution of the `op0` state inplace see [`evolutionOperatorIterator`](@ref).
The `diagonalize` argument decompose Hamiltonian into ``\\lambda_i, S, S^{-1}``
and calculate the exponential using eigenvalue decomposition. The `approximate` option asusmes that tspan is made of
equidistant points, therefore ``U(t)`` has to be calculated for ``t_0`` and ``t_\\text{step}``.

"""
function evolutionExact!(
    op_array::Array{Array{C,2},1},
    op0::Operator,
    tspan::Array,
    Hamiltonian::Operator,
) where {C<:ComputableType}
    N = length(tspan)
    buffer = zeros(C, size(op0.data))
    for t_i = 1:N
        push!(op_array, deepcopy(op0.data))
    end
    t_i = 0
    for U_op in evolutionOperatorIterator(Hamiltonian, tspan)
        t_i += 1
        buffer[:, :] .= convert(Array{C,2}, (U_op * op0 * U_op').data)
        op_array[t_i] = deepcopy(buffer)
        push!(op_array)
    end
    return op_array
end

"""
    evolutionApproximate(ket0, tspan, Hamiltonian)

Calculate approximate time evolution of the `ket0` based on ``U(t)`` that is
calculated for ``t_0`` and ``t_\\text{step}``. See [`evolutionOperator`](@ref).
This method returns Vector of Ket states.

"""
function evolutionApproximate(ket0::Ket, tspan::Array, Hamiltonian::Operator)
    N = length(tspan)
    ket_array = Array{typeof(ket0),1}(undef, 0)
    t_step = tspan[2] - tspan[1]
    U_op_step = evolutionOperator(Hamiltonian, t_step)
    U_op0 = evolutionOperator(Hamiltonian, tspan[1])
    ket = U_op0 * deepcopy(ket0)
    for t_i = 1:N
        push!(ket_array, deepcopy(ket))
    end
    for t_i = 2:N
        ket = U_op_step * ket
        ket_array[t_i] = deepcopy(ket)
    end
    return ket_array
end

"""
    evolutionApproximate!(ket_array, ket0, tspan, Hamiltonian)

Calculate approximate time evolution of the `ket0` inplace based on ``U(t)`` that is
calculated for ``t_0`` and ``t_\\text{step}``. See [`evolutionOperator`](@ref).
Argument `ket_array` is Vector of Arrays.

"""
function evolutionApproximate!(
    ket_array::Array{Array{C,1},1},
    ket0::Ket,
    tspan::Array,
    Hamiltonian::Operator,
) where {C<:ComputableType}
    N = length(tspan)
    t_step = tspan[2] - tspan[1]
    U_op_step = evolutionOperator(Hamiltonian, t_step)
    U_op0 = evolutionOperator(Hamiltonian, tspan[1])
    ket = U_op0 * deepcopy(ket0)
    buffer = zeros(C, size(ket.data))
    buffer[:] .= convert(Array{C,1}, ket.data)
    for t_i = 1:N
        push!(ket_array, deepcopy(buffer))
    end
    for t_i = 1:N-1
        ket = U_op_step * ket
        buffer[:] .= convert(Array{C,1}, ket.data)
        ket_array[t_i+1] = deepcopy(buffer)
    end
    return ket_array
end

"""
    evolutionApproximate(op0, tspan, Hamiltonian)

Calculate approximate time evolution of the `op0` based on ``U(t)`` that is
calculated for ``t_0`` and ``t_\\text{step}``. See [`evolutionOperator`](@ref).
This method returns Vector of Operators.

"""
function evolutionApproximate(op0::Operator, tspan::Array, Hamiltonian::Operator)
    N = length(tspan)
    op_array = Array{typeof(op0),1}(undef, 0)
    t_step = tspan[2] - tspan[1]
    U_op_step = evolutionOperator(Hamiltonian, t_step)
    U_op0 = evolutionOperator(Hamiltonian, tspan[1])
    U_op_step_d = U_op_step'
    op = U_op0 * deepcopy(op0) * U_op0'
    for t_i = 1:N
        push!(op_array, deepcopy(op))
    end
    for t_i = 2:N
        op = U_op_step * op * U_op_step_d
        op_array[t_i] = deepcopy(op)
    end
    return op_array
end

"""
    evolutionApproximate!(op_array, op0, tspan, Hamiltonian)

Calculate approximate time evolution of the `op0` inplace based on ``U(t)`` that is
calculated for ``t_0`` and ``t_\\text{step}``. See [`evolutionOperator`](@ref).
This method returns Vector of Arrays.

"""
function evolutionApproximate!(
    op_array::Array{Array{C,2},1},
    op0::Operator,
    tspan::Array,
    Hamiltonian::Operator,
) where {C<:ComputableType}
    N = length(tspan)
    t_step = tspan[2] - tspan[1]
    U_op_step = evolutionOperator(Hamiltonian, t_step)
    U_op0 = evolutionOperator(Hamiltonian, tspan[1])
    U_op_step_d = U_op_step'
    op = U_op0 * deepcopy(op0) * U_op0'
    buffer = zeros(C, size(op.data))
    buffer[:, :] .= convert(Array{C,2}, op.data)
    for t_i = 1:N
        push!(op_array, deepcopy(buffer))
    end
    for t_i = 2:N
        op = U_op_step * op * U_op_step_d
        buffer[:, :] .= convert(Array{C,2}, op.data)
        op_array[t_i] = deepcopy(buffer)
    end
    return op_array
end

"""
    evolution_exact(rho0, tspan, Hamiltonian; diagonalize = false)

Calculate exact time evolution of the `rho0` inplace based on ``U(t)``.
The `diagonalize` argument decompose Hamiltonian into ``\\lambda_i, S, S^{-1}``
and calculate the exponential using eigenvalue decomposition. See [`evolutionOperator`](@ref).
This method returns tspan and Vector of Operators.

"""
function evolution_exact(
    rho0::T,
    tspan::Array,
    Hamiltonian::U;
    diagonalize = false,
) where {B<:Basis,T<:Operator{B,B},U<:Operator{B,B}}
    N = length(tspan)

    rho = deepcopy(rho0)
    if diagonalize
        Ham_lambda, S = eigen(Hamiltonian.data)
        basis = GenericBasis([size(Ham_lambda, 1)])
        Sinv = inv(S)
        U_diagonal = zero(Ham_lambda)
        U_diagonal = map(lambda -> exp(-1im * lambda * tspan[1]), Ham_lambda)
        U_data = S * diagm(U_diagonal) * Sinv
        U_op = DenseOperator(basis, basis, U_data)
    else
        U_op = evolutionOperator(Hamiltonian, tspan[1])
    end
    rho = U_op * rho0 * U_op'
    rho_t = [rho]

    for t_i = 2:N
        t = tspan[t_i]
        if diagonalize
            U_diagonal .= map(lambda -> exp(-1im * lambda * t), Ham_lambda)
            U_op.data .= S * diagm(U_diagonal) * Sinv
        else
            U_op = evolutionOperator(Hamiltonian, t)
        end
        rho = U_op * rho0 * U_op'
        push!(rho_t, rho)
    end
    return tspan, rho_t
end

"""
    evolution_exact(rho0, tspan, Hamiltonian; diagonalize = false)

Calculate approximate time evolution of the `rho0` based on ``U(t)`` that is
calculated for ``t_0`` and ``t_\\text{step}``. The `diagonalize` argument decompose Hamiltonian into ``\\lambda_i, S, S^{-1}``
and calculate the exponential using eigenvalue decomposition. See [`evolutionOperator`](@ref).
This method returns tspan and Vector of Operators.

"""
function evolution_approximate(
    rho0::T,
    tspan::Array,
    Ham::U;
    diagonalize = false,
) where {B<:Basis,T<:Operator{B,B},U<:Operator{B,B}}
    N = length(tspan)
    t_step = tspan[2] - tspan[1]
    for t_i = 1:N-1
        if !(tspan[t_i+1] - tspan[t_i] ≈ t_step)
            throw(ErrorException("Points in in tspan have to be equidistant."))
        end
    end

    rho = deepcopy(rho0)
    if diagonalize
        Ham_lambda, S = eigen(Ham.data)
        basis = GenericBasis([size(Ham_lambda, 1)])
        Sinv = inv(S)

        U_diagonal = zero(Ham_lambda)
        U_diagonal = map(lambda -> exp(-1im * lambda * t_step), Ham_lambda)
        U_data = S * diagm(U_diagonal) * Sinv
        U_op_step = DenseOperator(basis, basis, U_data)

        U_diagonal = map(lambda -> exp(-1im * lambda * tspan[1]), Ham_lambda)
        U_data = S * diagm(U_diagonal) * Sinv
        U_op0 = DenseOperator(basis, basis, U_data)
    else
        U_op_step = evolutionOperator(Ham, t_step)
        U_op0 = evolutionOperator(Ham, tspan[1])
    end
    U_op_step_d = U_op_step'
    U_op0_d = U_op0'
    rho = U_op0 * rho0 * U_op0_d
    rho_t = [rho]

    for t_i = 2:N
        t = tspan[t_i]
        rho = U_op_step * rho * U_op_step_d
        push!(rho_t, rho)
    end
    return tspan, rho_t
end


text:
# module ithello

using ProximalAlgorithms.IterationTools
using Printf
using Base.Iterators

struct Index_iterable{I<:Int}
  x0::I     # initial point of the sequence
  m::I      # length of the sequence
  d::I      # number of repetitions of each sequence ex) 123123456... has d = 2
  N::I      # max of index
end

mutable struct Index_state{I<:Int}
  i::I      # current index: ex: 123123123456...
  m_n::I    # memory index behind
  m_o::I    # memory index front
  m::I      # length of the sequence
  cnt::I    # an inner counter
  x0::I     # an inner index
end

function Index_state(i::I, m_n::I, m_o::I, m::I) where {I}
    return Index_state{I}(
        i,
        m_n,
        m_o,
        m,
        0,
        1,

    )
end

function Base.iterate(iter::Index_iterable{I}) where {I}
    i =  (iter.x0 === nothing) ?  I(1) : iter.x0

    state = Index_state(i, I(1), I(2), iter.m)
    return state, state
end

function Base.iterate(
        iter::Index_iterable,
        state::Index_state,
    )

    state.cnt += 1

    if mod(state.cnt, iter.d * state.m) == 0
        state.x0 += state.m
        state.m_n = 0
        state.m_o = 1
        # the tail
        if state.x0 + state.m > iter.N
            state.m = iter.N - state.x0 + 1
            state.cnt = 0
            if state.m <= 0
                return nothing
            end
        end
    end
    # println(state.m)
    t = state.i
    state.i = mod(mod(t, iter.m), state.m) + state.x0

    t = state.m_n
    state.m_n = mod(t, state.m+1) + 1

    t = state.m_o
    state.m_o = mod(t, state.m+1) + 1

    return state, state
end

ext: jl


text:
using YaoQASM, Yao
using YaoQASM: U, CX, RESET, Barrier, MEASURE

using Test
@testset "yao api" begin
    # single qubit gate

    uh = U(5, 3, π/2, 0, π)
    @test mat(uh)*im ≈ mat(put(5, 3=>H))

    # cnot gate
    cn = CX(5, [3], [5])
    @test mat(cn) == mat(put(5, (3,5)=>ConstGate.CNOT))

    # reset gate
    rst = RESET(5, [3,5])
    @test rst isa Measure
    rst = RESET(5, 3)
    @test rst isa Measure

    @test Barrier{5}() isa Barrier

    ba = BitArray([false, false, false])
    fl = MEASURE(5, [3,4], view(ba, [2,3]))
    @test fl isa YaoQASM.FlushResults
    product_state(bit"11011") |> fl
    @test ba == BitArray([0,0,1])
end

@testset "YaoQASM.jl" begin
    src1 = """
    OPENQASM 2.0;

    gate cu1(theta) a,b
    {
        U(0,0,theta/2) a;
        CX a,b;
        U(0,0,-theta/2) b;
    }

    qreg q[3];
    qreg a[2];
    creg c[3];
    creg syn[2];
    cu1(pi/2) q[0],q[1];
    """
    @test_broken (string_ast ∘ parse_qasm ∘ lex)(src1) ==
"""YaoQASM.Grammar.Struct_mainprogram(
  ver=Token{real}(str=2.0, lineno=1, colno=1),
  prog=[
    YaoQASM.Grammar.Struct_gate(
      decl=YaoQASM.Grammar.Struct_gatedecl(
        id=Token{id}(str=cu1, lineno=3, colno=3),
        arglist1=Token{id}(str=theta, lineno=3, colno=3),
        arglist2=(
          Token{id}(str=a, lineno=3, colno=3),
          Token{id}(str=b, lineno=3, colno=3),
        ),
      ),
      goplist=[
        YaoQASM.Grammar.Struct_u(
          exprs=(
            (
              Token{nninteger}(str=0, lineno=5, colno=5),
              Token{nninteger}(str=0, lineno=5, colno=5),
            ),
            (
              Token{id}(str=theta, lineno=5, colno=5),
              Token{unnamed}(str=/, lineno=5, colno=5),
              Token{nninteger}(str=2, lineno=5, colno=5),
            ),
          ),
          arg=YaoQASM.Grammar.Struct_argument(
            id=Token{id}(str=a, lineno=5, colno=5),
            arg=nothing,
          ),
        ),
        YaoQASM.Grammar.Struct_cx(
          arg1=YaoQASM.Grammar.Struct_argument(
            id=Token{id}(str=a, lineno=6, colno=6),
            arg=nothing,
          ),
          arg2=YaoQASM.Grammar.Struct_argument(
            id=Token{id}(str=b, lineno=6, colno=6),
            arg=nothing,
          ),
        ),
        YaoQASM.Grammar.Struct_u(
          exprs=(
            (
              Token{nninteger}(str=0, lineno=7, colno=7),
              Token{nninteger}(str=0, lineno=7, colno=7),
            ),
            YaoQASM.Grammar.Struct_neg(
              value=(
                Token{id}(str=theta, lineno=7, colno=7),
                Token{unnamed}(str=/, lineno=7, colno=7),
                Token{nninteger}(str=2, lineno=7, colno=7),
              ),
            ),
          ),
          arg=YaoQASM.Grammar.Struct_argument(
            id=Token{id}(str=b, lineno=7, colno=7),
            arg=nothing,
          ),
        ),
      ],
    ),
    YaoQASM.Grammar.Struct_decl(
      regtype=Token{id}(str=qreg, lineno=10, colno=10),
      id=Token{id}(str=q, lineno=10, colno=10),
      int=Token{nninteger}(str=3, lineno=10, colno=10),
    ),
    YaoQASM.Grammar.Struct_decl(
      regtype=Token{id}(str=qreg, lineno=11, colno=11),
      id=Token{id}(str=a, lineno=11, colno=11),
      int=Token{nninteger}(str=2, lineno=11, colno=11),
    ),
    YaoQASM.Grammar.Struct_decl(
      regtype=Token{id}(str=creg, lineno=12, colno=12),
      id=Token{id}(str=c, lineno=12, colno=12),
      int=Token{nninteger}(str=3, lineno=12, colno=12),
    ),
    YaoQASM.Grammar.Struct_decl(
      regtype=Token{id}(str=creg, lineno=13, colno=13),
      id=Token{id}(str=syn, lineno=13, colno=13),
      int=Token{nninteger}(str=2, lineno=13, colno=13),
    ),
    YaoQASM.Grammar.Struct_iduop(
      op=Token{id}(str=cu1, lineno=14, colno=14),
      lst1=(
        Token{id}(str=pi, lineno=14, colno=14),
        Token{unnamed}(str=/, lineno=14, colno=14),
        Token{nninteger}(str=2, lineno=14, colno=14),
      ),
      lst2=(
        YaoQASM.Grammar.Struct_mixeditem(
          id=Token{id}(str=q, lineno=14, colno=14),
          arg=Token{nninteger}(str=0, lineno=14, colno=14),
        ),
        YaoQASM.Grammar.Struct_mixeditem(
          id=Token{id}(str=q, lineno=14, colno=14),
          arg=Token{nninteger}(str=1, lineno=14, colno=14),
        ),
      ),
    ),
  ],
)"""
    # Write your own tests here.
end

@testset "parse" begin
    src1 = """
    OPENQASM 2.0;

    gate cu1(theta) a,b
    {
        U(0,0,theta/2) a;
        CX a,b;
        U(0,0,-theta/2) b;
    }

    qreg q[3];
    qreg a[2];
    creg c[3];
    creg syn[2];
    cu1(pi/2) q[0],q[1];
    """
    circ = @eval $(trans(parse_qasm(lex(src1))))
    @show circ
    @test mat(circ) ≈ mat(
        chain(put(5, 1=>chain(rot(Z, π/4), rot(Y, 0.0), rot(Z, 0.0))),
            control(5, 1, 2=>X),
            put(5, 2=>chain(rot(Z, -π/4), rot(Y, 0.0), rot(Z, 0.0)))
            ))
    @test syn isa BitArray && length(syn) == 2
    @test c isa BitArray && length(c) == 3
    @test q isa Vector && length(q) == 3
    @test a isa Vector && length(a) == 2
end


text:
@enum StringEnum SE_FOO SE_BAR
@enum IntEnum IE_FOO IE_BAR

LightBSON.bson_representation_type(::Type{IntEnum}) = Int32

@testset "representations" begin

@testset "string enum" begin
    buf = empty!(fill(0xff, 1000))
    writer = BSONWriter(buf)
    writer["x"] = SE_BAR
    close(writer)
    @test BSONReader(buf, StrictBSONValidator())["x"][Any] == "SE_BAR"
    @test BSONReader(buf, StrictBSONValidator())["x"][String] == "SE_BAR"
    @test BSONReader(buf, StrictBSONValidator())["x"][StringEnum] == SE_BAR
end

@testset "int enum" begin
    buf = empty!(fill(0xff, 1000))
    writer = BSONWriter(buf)
    writer["x"] = IE_BAR
    close(writer)
    @test BSONReader(buf, StrictBSONValidator())["x"][Any] == Int32(IE_BAR)
    @test BSONReader(buf, StrictBSONValidator())["x"][Int32] == Int32(IE_BAR)
    @test BSONReader(buf, StrictBSONValidator())["x"][IntEnum] == IE_BAR
end

@testset "byte vector" begin
    buf = empty!(fill(0xff, 1000))
    writer = BSONWriter(buf)
    x = rand(UInt8, 10)
    writer["x"] = x
    close(writer)
    @test BSONReader(buf, StrictBSONValidator())["x"][BSONBinary].data == x
    @test BSONReader(buf, StrictBSONValidator())["x"][Vector{UInt8}] == x
    @test IndexedBSONReader(BSONIndex(10), BSONReader(buf, StrictBSONValidator()))["x"][BSONBinary].data == x
    @test IndexedBSONReader(BSONIndex(10), BSONReader(buf, StrictBSONValidator()))["x"][Vector{UInt8}] == x
end

@testset "tuple" begin
    buf = empty!(fill(0xff, 1000))
    writer = BSONWriter(buf)
    x = (1, 1.25, "test")
    writer["x"] = x
    close(writer)
    @test BSONReader(buf, StrictBSONValidator())["x"][Any] == collect(x)
    @test BSONReader(buf, StrictBSONValidator())["x"][typeof(x)] == x
end

@testset "IPv4" begin
    buf = empty!(fill(0xff, 1000))
    writer = BSONWriter(buf)
    x = ip"127.0.0.0"
    y = Sockets.InetAddr(x, 1234)
    writer["x"] = x
    writer["y"] = y
    close(writer)
    @test BSONReader(buf, StrictBSONValidator())["x"][String] == "127.0.0.0"
    @test BSONReader(buf, StrictBSONValidator())["x"][IPv4] == x
    @test BSONReader(buf, StrictBSONValidator())["y"][String] == "127.0.0.0:1234"
    @test BSONReader(buf, StrictBSONValidator())["y"][Sockets.InetAddr{IPv4}] == y
end

@testset "IPv6" begin
    buf = empty!(fill(0xff, 1000))
    writer = BSONWriter(buf)
    x = ip"::c01e:fc9a"
    y = Sockets.InetAddr(x, 1234)
    writer["x"] = x
    writer["y"] = y
    close(writer)
    @test BSONReader(buf, StrictBSONValidator())["x"][String] == "::c01e:fc9a"
    @test BSONReader(buf, StrictBSONValidator())["x"][IPv6] == x
    @test BSONReader(buf, StrictBSONValidator())["y"][String] == "[::c01e:fc9a]:1234"
    @test BSONReader(buf, StrictBSONValidator())["y"][Sockets.InetAddr{IPv6}] == y
end

end


text:
for F in (# :mod2pi, :square, :cube, :sqrt,
          :cbrt,
          :exp, :expm1, :log, :log1p, :log2, :log10,
          :sin, :cos, :tan, :csc, :sec, :cot,
          :asin, :acos, :atan, :acsc, :asec, :acot,
          :sinh, :cosh, :tanh, :csch, :sech, :coth,
          :asinh, :acosh, :atanh, :acsch, :asech, :acoth)
  @eval begin
    $F(x::Double32) = Double32($F(Double64(x)))
    $F(x::Double16) = Double16($F(Double64(x)))
    $F(x::DoubleFloat{Double32}) = DoubleFloat{Double32}($F(DoubleFloat{Double64}(x)))
    $F(x::DoubleFloat{Double16}) = DoubleFloat{Double16}($F(DoubleFloat{Double64}(x)))
  end
end

"""
    DomainGuard{T}

An easy way to provide mathematical functions
guards that disallow out-of-domain processing
describes the canonical value-span and offers
the manner of reductive domain comprehension.

For discontinous [gapping} periodic functions,
all that spans any of a subdomain collective
is included and admissable.  Each evaluand is
provisional; quantities with qualities map
the domain-reachable onto canonical values.

""" DomainGuard

#=
   missing values participate in exclusion
   given values bound regions of inclusion
=#

struct DomainGuard{T}

    maxneg::T           # most  negagtive magnitude of primary domain
    minneg::T           # least negative  magnitude of primary domain
    minpos::T           # zero or least positive magnitude of primary domain
    maxpos::T           # +inf or most positive magnitude of primary domain

    reduct::Function    # comprehensive domain >=>=> canonical domain

                        # boolean-valued states for each specialized numerica
    posinf::Bool
    posval::Bool
    negval::Bool
    neginf::Bool

    okzero::Bool

end


text:
module CayleyCrossover
using Random

export cayley_crossover!

# See MAGNA network alignment paper
# for each cycle in permutation
# cut off half off the cycle
function halfperm!(r, rng, visited = fill(false,length(r)))
    i = 1
    while i <= length(r)
        if visited[i]
            i += 1
            continue
        end
        # traverse cycle
        clen = 1
        cstart = i
        clast = cstart
        while true
            visited[clast] = true
            ccur = r[clast]
            if ccur == cstart; break; end
            clen += 1
            clast = ccur
        end
        cbreak = rand(rng,1:clen)
        # go part way into cycle
        for _ = 1:cbreak; clast = r[clast]; end
        # cut off rest of cycle and stitch it back
        cstart = clast
        clast = r[clast]
        for _ = 1:div(clen,2)
            ccur = r[clast]
            r[clast] = clast
            clast = ccur
        end
        r[cstart] = clast
        i += 1
    end
    r
end
function rdivperm!(r, p, q) # p q^-1
    for i = 1:length(p)
        r[q[i]] = p[i]
    end
    r
end

function cayley_crossover!(s::AbstractVector, p::AbstractVector, q::AbstractVector,
    rng::AbstractRNG,
    r::AbstractVector=similar(s),
    visited::AbstractVector=fill(false, length(s)))
    rdivperm!(r, p, q) # r = p q^-1
    halfperm!(r, rng, visited)
    # s[:] = r[q] # permute!(z, q) #
    @inbounds for i in 1:length(r)
        s[i] = r[q[i]]
    end
    s
end

end

module PermGA

using SparseArrays
using Random
using ..GAFramework
import ..GAFramework: fitness, crossover!, mutation!, selection, randcreature, printfitness
using ..CayleyCrossover

export PermCreature, NetalignModel

# Represents a permutation, i.e. 1-1 mapping from one set to another
struct PermCreature
    f :: Vector{Int}
    score :: Float64 # alignment score
end

fitness(x::PermCreature) where {T} = x.score

printfitness(curgen::Int, x::PermCreature) =
    println("curgen: $curgen value: $(x.f) score: $(x.score)")

"""
Model to find network alignment by optimizing S3 (see MAGNA paper)
A network alignment is represented using a permutation
Slow implementation but fine  for small graphs
    using GAFramework, GAFramework.PermGA
    using LightGraphs
    using Random
    mv1 = 8
    me1 = 15
    mv2 = 10
    me2 = 15
    G2 = adjacency_matrix(LightGraphs.erdos_renyi(mv2,me2))
    perm = randperm(mv2)
    G1 = G2[perm[1:mv1],perm[1:mv1]]
    # G1.|V| <= G2.|V|
    model = NetalignModel(G1,G2)
    st = GAState(model, ngen=100, npop=6_000, elite_fraction=0.1,
        params=Dict(:print_fitness_iter=>1))
    ga!(st)
    println("accuracy = ", sum(st.pop[1].f[1:mv1] .== perm[1:mv1])/mv1)

    G2 = sparse([1, 1, 2, 3, 4], [2, 3, 3, 4, 5], ones(Int,5), 5, 5)
    G2 = min.(1, (G2 + G2'))
    G1 = G2[1:4,1:4]
    model = NetalignModel(G1,G2)
    st = GAState(model, ngen=100, npop=6_000, elite_fraction=0.1,
        params=Dict(:print_fitness_iter=>1))
    ga!(st)
    println("accuracy = ", sum(st.pop[1].f[1:4] .== [1,2,3,4])/4)
    println("accuracy = ", sum(st.pop[1].f[1:4] .== [2,1,3,4])/4)
    st.pop[1], PermCreature(perm, model)
"""
struct NetalignModel <: GAModel
    G1   :: SparseMatrixCSC{Int,Int}
    G2   :: SparseMatrixCSC{Int,Int}
end

# aux = genauxga(..) can be used to reduce allocations here
# by using auxiliary space instead of reallocating
function PermCreature(f, m::NetalignModel)
    # Assumes that all edge weights are 1
    h = f[1:size(m.G1,1)] # view(f, 1:size(m.G1,1))
    K = m.G1 + m.G2[h,h]
    w = nonzeros(K)
    Nc = count(x->x==2, w)
    Nn = length(w) - Nc
    Nc,cr = divrem(Nc,2)
    Nn,nr = divrem(Nn,2) # need this & above so it works well with sim. ann. code
    if cr!=0 || nr!=0 error("G1 and G2 need to be symmetric"); end
    score = Nc/(Nc + Nn)
    PermCreature(f, score)
end

function randcreature(m::NetalignModel, aux, rng::AbstractRNG) where {T}
    f = randperm(rng, size(m.G2,1))
    PermCreature(f, m)
end

function crossover!(z, x, y,
    st::GAState{NetalignModel}, aux, rng::AbstractRNG) where {T}
    cayley_crossover!(z.f, x.f, y.f, rng)
    PermCreature(z.f, st.model)
end

end # PermGA


text:
### ============== ### ============== ### ============== ###
### OUTPUT FOLDER STRUCTURE
### ============== ### ============== ### ============== ###

function set_output_data_structure(path, N, κ, ω)

    parent_folder_path = "$(homedir())/art_DATA/$(path)"
    folder_path        = parent_folder_path * "/DATA/data_N_$(N)"
    reps_path          = folder_path * "/data_N_$(N)_k_$(κ)_w_$(ω)"

    try
        mkdir("$(homedir())/art_DATA")
    catch error
        println("Main data folder already exists")
    end

    try
        mkdir(parent_folder_path)
    catch error
        println("Parent folder already exists")
    end

    try
        mkdir(parent_folder_path * "/DATA")
    catch error
        println("Data folder already exists")
    end

    try
        mkdir(folder_path)
    catch error
        println("Folder already exists")
    end

    try
        mkdir(reps_path)
    catch error
        println("Parameter folder already exists")
    end

    return reps_path
end

### ============== ### ============== ### ============== ###

function set_output_data_structure(path, N, κ, ω, η)

    parent_folder_path = "$(homedir())/art_DATA/$(path)"
    folder_path        = parent_folder_path * "/DATA/data_N_$(N)"
    reps_path          = folder_path * "/data_N_$(N)_k_$(κ)_w_$(ω)_e_$(η)"

    try
        mkdir("$(homedir())/art_DATA")
    catch error
        println("Main data folder already exists")
    end

    try
        mkdir(parent_folder_path)
    catch error
        println("Parent folder already exists")
    end

    try
        mkdir(parent_folder_path * "/DATA")
    catch error
        println("Data folder already exists")
    end

    try
        mkdir(folder_path)
    catch error
        println("Folder already exists")
    end

    try
        mkdir(reps_path)
    catch error
        println("Parameter folder already exists")
    end

    return reps_path
end

### ============== ### ============== ### ============== ###
### COMPUTE RELATIVE (SQUARED) DISTANCES WITH METRIC THRESHOLD
### ============== ### ============== ### ============== ###

function calc_Rij_th(R_ij::SharedArray, pos::SharedArray, r0::Float64)

    @parallel for i in 1:3:length(pos)

        ri = div(i,3) + 1

        for j in (i+3):3:length(pos)

            rj = div(j,3) + 1

            d = (pos[i]-pos[j])^2 + (pos[i+1]-pos[j+1])^2 + (pos[i+2]-pos[j+2])^2
            d <= r0 ? R_ij[rj,ri] = 1 : R_ij[rj,ri] = -1
        end
    end
end

### ============== ### ============== ### ============== ###

function calc_Rij_th_2D(R_ij::SharedArray, pos::SharedArray, r0::Float64)

    @parallel for i in 1:2:length(pos)

        ri = div(i,2) + 1

        for j in (i+2):2:length(pos)

            rj = div(j,2) + 1

            d = (pos[i]-pos[j])^2 + (pos[i+1]-pos[j+1])^2
            d <= r0 ? R_ij[rj,ri] = 1 : R_ij[rj,ri] = -1
        end
    end
end

### ============== ### ============== ### ============== ###
### COMPUTE RAW RELATIVE (SQUARED) DISTANCES
### ============== ### ============== ### ============== ###

function calc_Rij(R_ij::SharedArray, pos::SharedArray)

    @parallel for i in 1:3:length(pos)

        ri = div(i,3) + 1

        for j in (i+3):3:length(pos)

            rj = div(j,3) + 1

            R_ij[rj,ri] = (pos[i]-pos[j])^2 + (pos[i+1]-pos[j+1])^2 + (pos[i+2]-pos[j+2])^2
        end
    end
end

### ============== ### ============== ### ============== ###

function calc_Rij_2D(R_ij::SharedArray, pos::SharedArray)

    @parallel for i in 1:2:length(pos)

        ri = div(i,2) + 1

        for j in (i+2):2:length(pos)

            rj = div(j,2) + 1

            R_ij[rj,ri] = (pos[i]-pos[j])^2 + (pos[i+1]-pos[j+1])^2
        end
    end

end

### ============== ### ============== ### ============== ###
### COMPUTE METRIC SHORT AND LONG RANGE INTERACTIONS
### ============== ### ============== ### ============== ###

function compute_metric_interactions(vel::SharedArray,v_r::SharedArray,v_n::SharedArray,R_ij::SharedArray,N::Int64, κ_dist)

    for id in first(localindexes(vel)):3:last(localindexes(vel))

        i = div(id, 3)
        # print(i+1,"|\t")

        v_r[3i+1] = 0.0
        v_r[3i+2] = 0.0
        v_r[3i+3] = 0.0

        v_n[3i+1] = 0.0
        v_n[3i+2] = 0.0
        v_n[3i+3] = 0.0

        sh_n = find(x -> x == 1, Symmetric(R_ij, :L)[(i*N)+1:(i+1)*N])
        ln_n = find(x -> x == -1, Symmetric(R_ij, :L)[(i*N)+1:(i+1)*N])

        k_sh = length(sh_n)

        # short-range
        if k_sh > 0
            for j in sh_n
                # print(j,"\t")
                v_r[3i+1] += vel[3(j-1)+1]
                v_r[3i+2] += vel[3(j-1)+2]
                v_r[3i+3] += vel[3(j-1)+3]
            end
            v_r[3i+1] /= k_sh
            v_r[3i+2] /= k_sh
            v_r[3i+3] /= k_sh
        end

        k_ln = rand(κ_dist)

        if k_ln > 0
            for j in sample(ln_n, k_ln, replace = true)
                v_n[3i+1] += vel[3(j-1)+1]
                v_n[3i+2] += vel[3(j-1)+2]
                v_n[3i+3] += vel[3(j-1)+3]
            end
            v_n[3i+1] /= k_ln
            v_n[3i+2] /= k_ln
            v_n[3i+3] /= k_ln
        end

    end
end

### ============== ### ============== ### ============== ###

function compute_metric_interactions_2D(vel::SharedArray,v_r::SharedArray,v_n::SharedArray,R_ij::SharedArray,N::Int64, κ_dist)

    for id in first(localindexes(vel)):2:last(localindexes(vel))

        i = div(id, 2)
        # print(i+1,"|\t")

        v_r[2i+1] = 0.0
        v_r[2i+2] = 0.0

        v_n[2i+1] = 0.0
        v_n[2i+2] = 0.0

        sh_n = find(x -> x == 1, Symmetric(R_ij, :L)[(i*N)+1:(i+1)*N])
        ln_n = find(x -> x == -1, Symmetric(R_ij, :L)[(i*N)+1:(i+1)*N])

        k_sh = length(sh_n)

        # short-range
        if k_sh > 0
            for j in sh_n
                # print(j,"\t")
                v_r[2i+1] += vel[2(j-1)+1]
                v_r[2i+2] += vel[2(j-1)+2]
            end
            v_r[2i+1] /= k_sh
            v_r[2i+2] /= k_sh
        end

        k_ln = rand(κ_dist)

        if k_ln > 0
            for j in sample(ln_n, k_ln, replace = true)
                v_n[2i+1] += vel[2(j-1)+1]
                v_n[2i+2] += vel[2(j-1)+2]
            end
            v_n[2i+1] /= k_ln
            v_n[2i+2] /= k_ln
        end

    end
end

### ============== ### ============== ### ============== ###
### COMPUTE TOPOLOGICAL SHORT AND LONG RANGE INTERACTIONS
### ============== ### ============== ### ============== ###

function compute_topological_interactions(vel::SharedArray,v_r::SharedArray,v_n::SharedArray,R_ij::SharedArray,N::Int64,k_sh::Int64, κ_dist)

    for id in first(localindexes(vel)):3:last(localindexes(vel))

        i = div(id, 3)

        v_r[3i+1] = 0.0
        v_r[3i+2] = 0.0
        v_r[3i+3] = 0.0

        v_n[3i+1] = 0.0
        v_n[3i+2] = 0.0
        v_n[3i+3] = 0.0

        neighbors = sortperm(Symmetric(R_ij, :L)[(i*N)+1:(i+1)*N])

        # short-range
        for j in neighbors[2:k_sh+1]
            v_r[3i+1] += vel[3(j-1)+1]
            v_r[3i+2] += vel[3(j-1)+2]
            v_r[3i+3] += vel[3(j-1)+3]
        end

        v_r[3i+1] /= k_sh
        v_r[3i+2] /= k_sh
        v_r[3i+3] /= k_sh

        k_ln = rand(κ_dist)

        # possible long range
        if k_ln > 0
            for j in sample(neighbors[k_sh+2:end], k_ln, replace = true)
                v_n[3i+1] += vel[3(j-1)+1]
                v_n[3i+2] += vel[3(j-1)+2]
                v_n[3i+3] += vel[3(j-1)+3]
            end
            v_n[3i+1] /= k_ln
            v_n[3i+2] /= k_ln
            v_n[3i+3] /= k_ln
        end

    end
end

### ============== ### ============== ### ============== ###

function compute_topological_interactions_2D(vel::SharedArray,v_r::SharedArray,v_n::SharedArray,R_ij::SharedArray,N::Int64,k_sh::Int64, κ_dist)

    for id in first(localindexes(vel)):3:last(localindexes(vel))

        i = div(id, 2)

        v_r[2i+1] = 0.0
        v_r[2i+2] = 0.0

        v_n[2i+1] = 0.0
        v_n[2i+2] = 0.0

        neighbors = sortperm(Symmetric(R_ij, :L)[(i*N)+1:(i+1)*N])

        # short-range
        for j in neighbors[2:k_sh+1]
            v_r[2i+1] += vel[2(j-1)+1]
            v_r[2i+2] += vel[2(j-1)+2]
        end

        v_r[2i+1] /= k_sh
        v_r[2i+2] /= k_sh

        k_ln = rand(κ_dist)

        # possible long range
        if k_ln > 0
            for j in sample(neighbors[k_sh+2:end], k_ln, replace = true)
                v_n[2i+1] += vel[2(j-1)+1]
                v_n[2i+2] += vel[2(j-1)+2]
            end
            v_n[2i+1] /= k_ln
            v_n[2i+2] /= k_ln
        end
    end
end

### ============== ### ============== ### ============== ###
### UPDATE PARTICLE'S POSITIONS AND VELOCITIES
### ============== ### ============== ### ============== ###

function update_particles(pos::SharedArray,vel::SharedArray,v_r::SharedArray,v_n::SharedArray,η::Float64,ω::Float64)

    q_r = Quaternion(zeros(Float64, 3))

    for id in first(localindexes(vel)):3:last(localindexes(vel))

        i = div(id, 3)

        signal = ω .* [v_r[3i+1] , v_r[3i+2], v_r[3i+3]] + (1.0 - ω) .* [v_n[3i+1], v_n[3i+2], v_n[3i+3]]

        p_vel = [vel[3i+1], vel[3i+2], vel[3i+3]]

        if norm(signal) > zero(Float64)

            signal_angle = dot(p_vel, signal) / (norm(signal)*norm(p_vel))

            signal_angle = ifelse( signal_angle < -1, -1, signal_angle)
            signal_angle = ifelse( signal_angle > 1, 1, signal_angle)

            q_r = qrotation(cross(p_vel, signal), acos(signal_angle) + η * (2.0 * rand() * pi - pi)) * Quaternion(p_vel)

        else

            noise = randn(3)

            q_r = qrotation(cross(p_vel, noise), η * (2.0 * rand() * pi - pi)) * Quaternion(p_vel)

        end

        u_vel = normalize([q_r.v1, q_r.v2, q_r.v3])

        vel[3i+1] = u_vel[1]
        vel[3i+2] = u_vel[2]
        vel[3i+3] = u_vel[3]

        pos[3i+1] += u_vel[1]
        pos[3i+2] += u_vel[2]
        pos[3i+3] += u_vel[3]

    end
end

### ============== ### ============== ### ============== ###

function update_particles_2D(pos::SharedArray,vel::SharedArray,v_r::SharedArray,v_n::SharedArray,η::Float64,ω::Float64)

    for id in first(localindexes(vel)):2:last(localindexes(vel))

        i = div(id, 2)

        signal = ω .* [v_r[2i+1] , v_r[2i+2]] + (1.0 - ω) .* [v_n[2i+1], v_n[2i+2]]

        p_vel = [vel[2i+1], vel[2i+2]]
        n_vel = zeros(Float64, 2)

        if norm(signal) > zero(Float64)

            signal_angle = dot(p_vel, signal) / (norm(signal)*norm(p_vel))

            signal_angle = ifelse( signal_angle < -1, -1, signal_angle)
            signal_angle = ifelse( signal_angle > 1, 1, signal_angle)

            c = cos(signal_angle + η * (2.0 * rand() * pi - pi))
            s = sin(signal_angle + η * (2.0 * rand() * pi - pi))

            n_vel[1] = p_vel[1]*c - p_vel[2]*s;
            n_vel[2] = p_vel[1]*s + p_vel[2]*c;

        else

            c = cos(η * (2.0 * rand() * pi - pi))
            s = sin(η * (2.0 * rand() * pi - pi))

            n_vel[1] = p_vel[1]*c - p_vel[2]*s;
            n_vel[2] = p_vel[1]*s + p_vel[2]*c;

        end

        normalize!(n_vel)

        vel[2i+1] = n_vel[1]
        vel[2i+2] = n_vel[2]

        pos[2i+1] += n_vel[1]
        pos[2i+2] += n_vel[2]

    end
end

### ============== ### ============== ### ============== ###
### SYSTEM UPDATE (METRIC SHORT-RANGE INTERACTIONS)
### ============== ### ============== ### ============== ###

function evolve_metric_system(pos::SharedArray, vel::SharedArray, v_r::SharedArray, v_n::SharedArray, R_ij::SharedArray, r0::Float64, N::Int64,η::Float64,ω::Float64,κ_dist)

    calc_Rij_th(R_ij, pos, r0)

    @sync begin
        for p in workers()
            @async remotecall_wait(compute_metric_interactions, p, vel, v_r, v_n, R_ij, N, κ_dist)
        end
    end

    @sync begin
        for p in workers()
            @async remotecall_wait(update_particles, p, pos, vel, v_r, v_n, η, ω)
        end
    end

end

### ============== ### ============== ### ============== ###

function evolve_metric_system_2D(pos::SharedArray, vel::SharedArray, v_r::SharedArray, v_n::SharedArray, R_ij::SharedArray, r0::Float64, N::Int64,η::Float64,ω::Float64,κ_dist)

    calc_Rij_th(R_ij, pos, r0)

    @sync begin
        for p in workers()
            @async remotecall_wait(compute_metric_interactions_2D, p, vel, v_r, v_n, R_ij, N, κ_dist)
        end
    end

    @sync begin
        for p in workers()
            @async remotecall_wait(update_particles_2D, p, pos, vel, v_r, v_n, η, ω)
        end
    end

end

### ============== ### ============== ### ============== ###
### SYSTEM UPDATE (TOPOLOGICAL SHORT-RANGE INTERACTIONS)
### ============== ### ============== ### ============== ###

function evolve_topological_system(pos::SharedArray, vel::SharedArray, v_r::SharedArray, v_n::SharedArray, R_ij::SharedArray,N::Int64,k_sh::Int64,η::Float64,ω::Float64,κ_dist)

    calc_Rij(R_ij, pos)

    @sync begin
        for p in workers()
            @async remotecall_wait(compute_topological_interactions, p, vel, v_r, v_n, R_ij, N, k_sh, κ_dist)
        end
    end

    @sync begin
        for p in workers()
            @async remotecall_wait(update_particles, p, pos, vel, v_r, v_n, η, ω)
        end
    end

end

### ============== ### ============== ### ============== ###

function evolve_topological_system_2D(pos::SharedArray, vel::SharedArray, v_r::SharedArray, v_n::SharedArray, R_ij::SharedArray,N::Int64,k_sh::Int64,η::Float64,ω::Float64,κ_dist)

    calc_Rij(R_ij, pos)

    @sync begin
        for p in workers()
            @async remotecall_wait(compute_topological_interactions_2D, p, vel, v_r, v_n, R_ij, N, k_sh, κ_dist)
        end
    end

    @sync begin
        for p in workers()
            @async remotecall_wait(update_particles_2D, p, pos, vel, v_r, v_n, η, ω)
        end
    end

end


text:
# Run molecular dynamics
# See https://www.saylor.org/site/wp-content/uploads/2011/06/MA221-6.1.pdf for
#   integration algorithm - used shorter second version
# See https://udel.edu/~arthij/MD.pdf for information on forces
# See https://arxiv.org/pdf/1401.1181.pdf for applying forces to atoms
# See Gromacs manual for other aspects of forces

export
    temperature,
    simulate!

"The constant for Coulomb interaction, 1/(4*π*ϵ0*ϵr)."
const coulomb_const = 138.935458 / 70.0 # Treat ϵr as 70 for now

"Square of the non-bonded interaction distance cutoff in nm^2."
const sqdist_cutoff_nb = 1.0 ^ 2

"Square of the neighbour list distance cutoff in nm^2."
const sqdist_cutoff_nl = 1.5 ^ 2

"3D acceleration values, e.g. for an atom, in nm/(ps^2)."
mutable struct Acceleration <: FieldVector{3, Float64}
    x::Float64
    y::Float64
    z::Float64
end

"Update coordinates of all atoms and bound to the bounding box."
function update_coordinates!(coords::Vector{Coordinates},
                    velocities::Vector{Velocity},
                    accels::Vector{Acceleration},
                    timestep::Real,
                    box_size::Real)
    for (i, c) in enumerate(coords)
        c.x += velocities[i].x*timestep + 0.5*accels[i].x*timestep^2
        while (c.x >= box_size) c.x -= box_size end
        while (c.x < 0.0) c.x += box_size end

        c.y += velocities[i].y*timestep + 0.5*accels[i].y*timestep^2
        while (c.y >= box_size) c.y -= box_size end
        while (c.y < 0.0) c.y += box_size end

        c.z += velocities[i].z*timestep + 0.5*accels[i].z*timestep^2
        while (c.z >= box_size) c.z -= box_size end
        while (c.z < 0.0) c.z += box_size end
    end
    return coords
end

"Update velocities of all atoms using the accelerations."
function update_velocities!(velocities::Vector{Velocity},
                    accels_t::Vector{Acceleration},
                    accels_t_dt::Vector{Acceleration},
                    timestep::Real)
    for (i, v) in enumerate(velocities)
        v.x += 0.5*(accels_t[i].x+accels_t_dt[i].x)*timestep
        v.y += 0.5*(accels_t[i].y+accels_t_dt[i].y)*timestep
        v.z += 0.5*(accels_t[i].z+accels_t_dt[i].z)*timestep
    end
    return velocities
end

"Rescale random velocities according to the Andersen thermostat."
function rescale_velocities!(universe::Universe,
                    timestep::Real;
                    coupling_const::Real=0.1)
    T = universe.temperature
    for (i, v) in enumerate(universe.velocities)
        if rand() < timestep / coupling_const
            mass = universe.molecule.atoms[i].mass
            v.x = maxwellboltzmann(mass, T)
            v.y = maxwellboltzmann(mass, T)
            v.z = maxwellboltzmann(mass, T)
        end
    end
end

"Vector between two coordinate values, accounting for the bounding box."
function vector1D(c1::Real, c2::Real, box_size::Real)
    if c1 < c2
        return (c2 - c1) < (c1 - c2 + box_size) ? (c2 - c1) : (c2 - c1 - box_size)
    else
        return (c1 - c2) < (c2 - c1 + box_size) ? (c2 - c1) : (c2 - c1 + box_size)
    end
end

"3D vector between two `Coordinates`, accounting for the bounding box."
vector(coords_one::Coordinates, coords_two::Coordinates, box_size::Real) = [
        vector1D(coords_one.x, coords_two.x, box_size),
        vector1D(coords_one.y, coords_two.y, box_size),
        vector1D(coords_one.z, coords_two.z, box_size)]

"Square distance between two `Coordinates`, accounting for the bounding box."
sqdist(coords_one::Coordinates, coords_two::Coordinates, box_size::Real) =
        vector1D(coords_one.x, coords_two.x, box_size) ^ 2 +
        vector1D(coords_one.y, coords_two.y, box_size) ^ 2 +
        vector1D(coords_one.z, coords_two.z, box_size) ^ 2

"Force on each atom in a covalent bond."
function forcebond(coords_one::Coordinates,
                coords_two::Coordinates,
                bondtype::Bondtype,
                box_size::Real)
    ab = vector(coords_one, coords_two, box_size)
    c = bondtype.kb * (norm(ab) - bondtype.b0)
    fa = c*normalize(ab)
    fb = -fa
    return fa..., fb...
end

# Sometimes domain error occurs for acos if the float is > 1.0 or < 1.0
acosbound(x::Real) = acos(max(min(x, 1.0), -1.0))

"Force on each atom in a bond angle."
function forceangle(coords_one::Coordinates,
                coords_two::Coordinates,
                coords_three::Coordinates,
                angletype::Angletype,
                box_size::Real)
    ba = vector(coords_two, coords_one, box_size)
    bc = vector(coords_two, coords_three, box_size)
    pa = normalize(ba × (ba × bc))
    pc = normalize(-bc × (ba × bc))
    angle_term = -angletype.cth * (acosbound(dot(ba, bc) / (norm(ba) * norm(bc))) - angletype.th0)
    fa = (angle_term / norm(ba)) * pa
    fc = (angle_term / norm(bc)) * pc
    fb = -fa - fc
    return fa..., fb..., fc...
end

"Force on each atom in a dihedral torsion angle."
function forcedihedral(coords_one::Coordinates,
                coords_two::Coordinates,
                coords_three::Coordinates,
                coords_four::Coordinates,
                dihedraltype::Dihedraltype,
                box_size::Real)
    ba = vector(coords_two, coords_one, box_size)
    bc = vector(coords_two, coords_three, box_size)
    dc = vector(coords_four, coords_three, box_size)
    p1 = normalize(ba × bc)
    p2 = normalize(-dc × -bc)
    θ = atan(dot((-ba × bc) × (bc × -dc), normalize(bc)), dot(-ba × bc, bc × -dc))
    angle_term = 0.5*(dihedraltype.f1*sin(θ) - 2*dihedraltype.f2*sin(2*θ) + 3*dihedraltype.f3*sin(3*θ))
    fa = (angle_term / (norm(ba) * sin(acosbound(dot(ba, bc) / (norm(ba) * norm(bc)))))) * p1
    fd = (angle_term / (norm(dc) * sin(acosbound(dot(bc, dc) / (norm(bc) * norm(dc)))))) * p2
    oc = 0.5 * bc
    tc = -(oc × fd + 0.5 * (-dc × fd) + 0.5 * (ba × fa))
    fc = (1 / dot(oc, oc)) * (tc × oc)
    fb = -fa - fc -fd
    return fa..., fb..., fc..., fd...
end

"Force on each atom due to Lennard Jones attractive/repulsive potential."
function forcelennardjones end

@fastmath @inbounds function forcelennardjones(coords_one::Coordinates,
                coords_two::Coordinates,
                atom_one::Atom,
                atom_two::Atom,
                box_size::Real)
    σ = sqrt(atom_one.σ * atom_two.σ)
    ϵ = sqrt(atom_one.ϵ * atom_two.ϵ)
    dx = vector1D(coords_one.x, coords_two.x, box_size)
    dy = vector1D(coords_one.y, coords_two.y, box_size)
    dz = vector1D(coords_one.z, coords_two.z, box_size)
    r2 = dx*dx + dy*dy + dz*dz
    if r2 > sqdist_cutoff_nb
        return 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
    end
    six_term = (σ^2 / r2)^3
    # Limit this to 100 as a fudge to stop it exploding
    f = min(((24 * ϵ) / (r2)) * (2 * six_term^2 - six_term), 100.0)
    return -f*dx, -f*dy, -f*dz, f*dx, f*dy, f*dz
end

"Force on each atom due to electrostatic Coulomb potential."
function forcecoulomb end

@fastmath @inbounds function forcecoulomb(coords_one::Coordinates,
                coords_two::Coordinates,
                charge_one::Real,
                charge_two::Real,
                box_size::Real)
    dx = vector1D(coords_one.x, coords_two.x, box_size)
    dy = vector1D(coords_one.y, coords_two.y, box_size)
    dz = vector1D(coords_one.z, coords_two.z, box_size)
    r2 = dx*dx + dy*dy + dz*dz
    if r2 > sqdist_cutoff_nb
        return 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
    end
    f = (coulomb_const * charge_one * charge_two) / sqrt(r2^3)
    return -f*dx, -f*dy, -f*dz, f*dx, f*dy, f*dz
end

"Update accelerations of all atoms using the bonded and non-bonded forces."
function update_accelerations!(accels::Vector{Acceleration},
                universe::Universe,
                forcefield::Forcefield)
    # Clear accelerations
    for a in accels
        fill!(a, 0.0)
    end

    # Bonded forces
    # Covalent bond forces
    #temp_atomid = 59
    #temp_bond_sum = 0.0
    atoms = universe.molecule.atoms
    for b in universe.molecule.bonds
        if haskey(forcefield.bondtypes, "$(atoms[b.atom_i].attype)/$(atoms[b.atom_j].attype)")
            bondtype = forcefield.bondtypes["$(atoms[b.atom_i].attype)/$(atoms[b.atom_j].attype)"]
        else
            bondtype = forcefield.bondtypes["$(atoms[b.atom_j].attype)/$(atoms[b.atom_i].attype)"]
        end
        d1x, d1y, d1z, d2x, d2y, d2z = forcebond(
            universe.coords[b.atom_i], universe.coords[b.atom_j],
            bondtype, universe.box_size)
        accels[b.atom_i].x += d1x
        accels[b.atom_i].y += d1y
        accels[b.atom_i].z += d1z
        accels[b.atom_j].x += d2x
        accels[b.atom_j].y += d2y
        accels[b.atom_j].z += d2z
        #b.atom_i == temp_atomid ? temp_bond_sum += d1x : nothing
        #b.atom_j == temp_atomid ? temp_bond_sum += d2x : nothing
    end
    #print("Bond: $(round(temp_bond_sum, 2)), ")

    # Angle forces
    #temp_angle_sum = 0.0
    for a in universe.molecule.angles
        if haskey(forcefield.angletypes, "$(atoms[a.atom_i].attype)/$(atoms[a.atom_j].attype)/$(atoms[a.atom_k].attype)")
            angletype = forcefield.angletypes["$(atoms[a.atom_i].attype)/$(atoms[a.atom_j].attype)/$(atoms[a.atom_k].attype)"]
        else
            angletype = forcefield.angletypes["$(atoms[a.atom_k].attype)/$(atoms[a.atom_j].attype)/$(atoms[a.atom_i].attype)"]
        end
        d1x, d1y, d1z, d2x, d2y, d2z, d3x, d3y, d3z = forceangle(
            universe.coords[a.atom_i], universe.coords[a.atom_j],
            universe.coords[a.atom_k], angletype, universe.box_size)
        accels[a.atom_i].x += d1x
        accels[a.atom_i].y += d1y
        accels[a.atom_i].z += d1z
        accels[a.atom_j].x += d2x
        accels[a.atom_j].y += d2y
        accels[a.atom_j].z += d2z
        accels[a.atom_k].x += d3x
        accels[a.atom_k].y += d3y
        accels[a.atom_k].z += d3z
        #a.atom_i == temp_atomid ? temp_angle_sum += d1x : nothing
        #a.atom_j == temp_atomid ? temp_angle_sum += d2x : nothing
        #a.atom_k == temp_atomid ? temp_angle_sum += d3x : nothing
    end
    #print("Angle: $(round(temp_angle_sum, 2)), ")

    # Dihedral forces
    #temp_tor_sum = 0.0
    for d in universe.molecule.dihedrals
        dihedraltype = forcefield.dihedraltypes["$(atoms[d.atom_i].attype)/$(atoms[d.atom_j].attype)/$(atoms[d.atom_k].attype)/$(atoms[d.atom_l].attype)"]
        d1x, d1y, d1z, d2x, d2y, d2z, d3x, d3y, d3z, d4x, d4y, d4z = forcedihedral(
            universe.coords[d.atom_i], universe.coords[d.atom_j],
            universe.coords[d.atom_k], universe.coords[d.atom_l],
            dihedraltype, universe.box_size)
        accels[d.atom_i].x += d1x
        accels[d.atom_i].y += d1y
        accels[d.atom_i].z += d1z
        accels[d.atom_j].x += d2x
        accels[d.atom_j].y += d2y
        accels[d.atom_j].z += d2z
        accels[d.atom_k].x += d3x
        accels[d.atom_k].y += d3y
        accels[d.atom_k].z += d3z
        accels[d.atom_l].x += d4x
        accels[d.atom_l].y += d4y
        accels[d.atom_l].z += d4z
        #d.atom_i == temp_atomid ? temp_tor_sum += d1x : nothing
        #d.atom_j == temp_atomid ? temp_tor_sum += d2x : nothing
        #d.atom_k == temp_atomid ? temp_tor_sum += d3x : nothing
        #d.atom_l == temp_atomid ? temp_tor_sum += d4x : nothing
    end
    #print("Torsion: $(round(temp_tor_sum, 2)), ")

    # Non-bonded forces
    #temp_lj_sum = 0.0
    #temp_el_sum = 0.0
    for (i, j, d) in universe.neighbour_list
        a1 = universe.molecule.atoms[i]
        a2 = universe.molecule.atoms[j]

        # Lennard Jones forces
        if a1.σ > 0.0 && a2.σ > 0.0
            d1x, d1y, d1z, d2x, d2y, d2z = forcelennardjones(
                universe.coords[i], universe.coords[j], a1, a2, universe.box_size)
            # Non-bonded forces are halved if the pair is a specified non-bonded pair
            accels[i].x += d ? 0.5*d1x : d1x
            accels[i].y += d ? 0.5*d1y : d1y
            accels[i].z += d ? 0.5*d1z : d1z
            accels[j].x += d ? 0.5*d2x : d2x
            accels[j].y += d ? 0.5*d2y : d2y
            accels[j].z += d ? 0.5*d2z : d2z
            #i == temp_atomid ? temp_lj_sum += d1x : nothing
            #j == temp_atomid ? temp_lj_sum += d2x : nothing
        end

        # Electrostatic forces
        d1x, d1y, d1z, d2x, d2y, d2z = forcecoulomb(
            universe.coords[i], universe.coords[j], a1.charge, a2.charge, universe.box_size)
        accels[i].x += d ? 0.5*d1x : d1x
        accels[i].y += d ? 0.5*d1y : d1y
        accels[i].z += d ? 0.5*d1z : d1z
        accels[j].x += d ? 0.5*d2x : d2x
        accels[j].y += d ? 0.5*d2y : d2y
        accels[j].z += d ? 0.5*d2z : d2z
        #i == temp_atomid ? temp_el_sum += d1x : nothing
        #j == temp_atomid ? temp_el_sum += d2x : nothing
    end
    #print("LJ: $(round(temp_lj_sum, 2)), ")
    #print("Coulomb: $(round(temp_el_sum, 2)), ")

    # Divide sum of forces my the atom mass to get acceleration
    for i in 1:length(accels)
        accels[i].x /= universe.molecule.atoms[i].mass
        accels[i].y /= universe.molecule.atoms[i].mass
        accels[i].z /= universe.molecule.atoms[i].mass
    end

    return accels
end

"Update list of close atoms between which non-bonded forces are calculated."
function update_neighbours!(universe::Universe)
    empty!(universe.neighbour_list)
    for i in 1:length(universe.coords)
        ci = universe.coords[i]
        nbmi = universe.molecule.nb_matrix[:, i]
        nbpi = universe.molecule.nb_pairs[:, i]
        for j in 1:(i-1)
            if sqdist(ci, universe.coords[j], universe.box_size) < sqdist_cutoff_nl && nbmi[j]
                push!(universe.neighbour_list, (i, j, nbpi[j]))
            end
        end
    end
    return universe
end

"Calculate the temperature of a system from the kinetic energy of the atoms."
function temperature(u::Universe)
    ke = sum([a.mass * dot(u.velocities[i], u.velocities[i]) for (i, a) in enumerate(u.molecule.atoms)]) / 2
    df = 3*length(u.velocities) - 3
    return 2*ke / df
end

"Initialise empty `Acceleration`s."
empty_accelerations(n_atoms::Integer) = [Acceleration(0.0, 0.0, 0.0) for i in 1:n_atoms]

"Simulate molecular dynamics."
function simulate!(s::Simulation, n_steps::Integer)
    report("Starting simulation")
    n_atoms = length(s.universe.coords)
    update_neighbours!(s.universe)
    a_t = update_accelerations!(empty_accelerations(n_atoms), s.universe, s.forcefield)
    a_t_dt = empty_accelerations(n_atoms)
    out_prefix = "pdbs_5XER"
    #writepdb("$out_prefix/snap_0.pdb", s.universe)
    @showprogress for i in 1:n_steps
        update_coordinates!(s.universe.coords, s.universe.velocities, a_t,
                            s.timestep, s.universe.box_size)
        update_accelerations!(a_t_dt, s.universe, s.forcefield)
        update_velocities!(s.universe.velocities, a_t, a_t_dt, s.timestep)
        rescale_velocities!(s.universe, s.timestep)
        # Update neighbour list every 10 steps
        if i % 10 == 0
            update_neighbours!(s.universe)
            #writepdb("$out_prefix/snap_$(Int(i/10)).pdb", s.universe)
        end
        push!(s.temperatures, temperature(s.universe))
        a_t = a_t_dt
        s.steps_made += 1
    end
    report("Simulation finished")
    return s
end

simulate!(s::Simulation) = simulate!(s, s.n_steps - s.steps_made)


text:
using PyTensorFlow
using PyCall
using LinearAlgebra
using PyPlot
using Random
Random.seed!(233)

if Sys.islinux()
py"""
import tensorflow as tf
libUpwlapOp = tf.load_op_library('build/libUpwlapOp.so')
@tf.custom_gradient
def upwlap_op(perm,mobi,func,h,rhograv):
    out = libUpwlapOp.upwlap_op(perm,mobi,func,h,rhograv)
    def grad(dy):
        return libUpwlapOp.upwlap_op_grad(dy, out, perm,mobi,func,h,rhograv)
    return out, grad
"""
elseif Sys.isapple()
py"""
import tensorflow as tf
libUpwlapOp = tf.load_op_library('build/libUpwlapOp.dylib')
@tf.custom_gradient
def upwlap_op(perm,mobi,func,h,rhograv):
    out = libUpwlapOp.upwlap_op(perm,mobi,func,h,rhograv)
    def grad(dy):
        return libUpwlapOp.upwlap_op_grad(dy, out, perm,mobi,func,h,rhograv)
    return out, grad
"""
elseif Sys.iswindows()
py"""
import tensorflow as tf
libUpwlapOp = tf.load_op_library('build/libUpwlapOp.dll')
@tf.custom_gradient
def upwlap_op(perm,mobi,func,h,rhograv):
    out = libUpwlapOp.upwlap_op(perm,mobi,func,h,rhograv)
    def grad(dy):
        return libUpwlapOp.upwlap_op_grad(dy, out, perm,mobi,func,h,rhograv)
    return out, grad
"""
end

upwlap_op = py"upwlap_op"

# TODO:
# u = upwlap_op(perm,mobi,func,h,rhograv)
# sess = Session()
# init(sess)
# run(sess, u)

# TODO:
h = 1.0
rho = 1000.0
G = 9.8
len_z = 16
len_x = 32
nz = Int(len_z/h + 1)
nx = Int(len_x/h + 1)
tf_h=constant(1.0)

perm = rand(nz, nx)
mobi = rand(nz, nx)
func = rand(nz, nx)

tf_perm = constant(perm)
tf_mobi = constant(mobi)
tf_func = constant(func)

# gradient check -- v
function scalar_function(m)
    # return sum(tanh(upwlap_op(m, tf_mobi, tf_func, tf_h, constant(rho*G))))
    # return sum(tanh(upwlap_op(tf_perm, m, tf_func, tf_h, constant(rho*G))))
    return sum(tanh(upwlap_op(tf_perm, tf_mobi, m, tf_h, constant(rho*G))))
end

# m_ = constant(rand(10,20))
# m_ = tf_perm
# m_ = tf_mobi
m_ = tf_func
v_ = rand(nz, nx)
y_ = scalar_function(m_)
dy_ = gradients(y_, m_)
ms_ = Array{Any}(undef, 5)
ys_ = Array{Any}(undef, 5)
s_ = Array{Any}(undef, 5)
w_ = Array{Any}(undef, 5)
gs_ =  @. 1 / 20^(1:5)

for i = 1:5
    g_ = gs_[i]
    ms_[i] = m_ + g_*v_
    ys_[i] = scalar_function(ms_[i])
    s_[i] = ys_[i] - y_
    w_[i] = s_[i] - g_*sum(v_.*dy_)
end

sess = Session()
init(sess)
sval_ = run(sess, s_)
wval_ = run(sess, w_)
close("all")
loglog(gs_, abs.(sval_), "*-", label="finite difference")
loglog(gs_, abs.(wval_), "+-", label="automatic differentiation")
loglog(gs_, gs_.^2 * 0.5*abs(wval_[1])/gs_[1]^2, "--",label="\$\\mathcal{O}(\\gamma^2)\$")
loglog(gs_, gs_ * 0.5*abs(sval_[1])/gs_[1], "--",label="\$\\mathcal{O}(\\gamma)\$")

plt.gca().invert_xaxis()
legend()
xlabel("\$\\gamma\$")
ylabel("Error")


text:
import DataStructures: DefaultDict, OrderedSet
import LinearAlgebra: cross, norm

function match(segments::Vector{LineSegment})::Vector{Polygon}
    cache = DefaultDict{Vertex, Vector{Vertex}}(Vector{Vertex})
    for segment in segments
        push!(cache[segment.a], segment.b)
        push!(cache[segment.b], segment.a)
    end

    polygons = Vector{Polygon}()

    while !isempty(cache)
        polygon = OrderedSet{Vertex}()
        start = first(cache).first
        push!(polygon, start)
        current = pop!(cache[start])
        delete!(cache, start)

        while haskey(cache, current)
            push!(polygon, current)
            a = pop!(cache[current])
            b = pop!(cache[current])
            @assert isempty(cache[current])
            delete!(cache, current)
            current = in(a, polygon) ? b : a
            push!(polygon, a)
            push!(polygon, b)
        end

        push!(polygons, collect(polygon))
    end

    return polygons
end

function simplify(polygons::Vector{Polygon})
    map(simplify, polygons)
end

function simplify(polygon::Polygon)
    [b for ((a, b), (c, d)) in edges(collect(edges(polygon))) if !iscollinear(b - a, d - c)]
end

function iscollinear(a::Vertex, b::Vertex)
    isapprox(norm(cross(a, b)), 0, atol=1e-2)
end

function correct(polygons::Vector{Polygon})
    corrected = Vector{Polygon}()
    boxes = map(boundingbox, polygons)
    for (i, box) in enumerate(boxes)
        num_enclosing = sum(map(other -> issubset(box, other) ? 1 : 0, boxes))
        if iseven(num_enclosing)
            push!(corrected, isclockwise(polygons[i]) ? reverse(polygons[i]) : polygons[i])
        else
            push!(corrected, !isclockwise(polygons[i]) ? reverse(polygons[i]) : polygons[i])
        end
    end
    corrected
end


text:
import Pkg; Pkg.add("Documenter")
push!(LOAD_PATH,"../src/")
using Documenter, XMLParser
makedocs(
	sitename = "XMLParser.jl",
	modules = [XMLParser],
	pages = [
		"Home" => "index.md"
	]
	)
deploydocs(
    repo = "github.com/baxmittens/XMLParser.jl.git"
)

max_stars_repo_name: baxmittens/XMLParser

max_issues_repo_name: baxmittens/XMLParser

max_forks_repo_name: baxmittens/XMLParser


text:
function discrete_adjoint(sol, dg, t; autojacvec=ForwardDiffVJP(), kwargs...)

    @assert sol.dense "Currently `discrete_adjoint` only works with dense solutions"
    @assert all(tidx -> tidx in sol.t, t) "A `tstop` must be set for every data point"

    @unpack prob, alg = sol
    @unpack f, p, u0, tspan = prob
    t0, tf = tspan

    integrator = OrdinaryDiffEq.init(prob, alg; kwargs...)

    # state jacobian function
    jac = state_jacobian(integrator)

    # vector-jacobian products
    vjp = vector_jacobian_product_function(integrator, autojacvec)

    # Discrete Adjoint Solution
    dp = zeros(length(p))
    du0 = zeros(length(u0))
    λ = zeros(length(u0))
    rhs = zeros(length(u0))
    resid = zeros(length(u0))
    J = zeros(length(u0), length(u0))
    dgval = zeros(length(u0))
    uvjpval = zeros(length(u0))
    pvjpval = zeros(length(p))
    for i = length(sol):-1:2
        # arguments for this step
        ti = sol.t[i]
        dt = ti - sol.t[i-1]
        uprev = sol.u[i-1]
        ui = sol.u[i]
        # compute right hand side
        idx = findfirst(tidx -> tidx == ti, t)
        if !isnothing(idx)
            dg(dgval, ui, p, ti, idx)
            rhs .-= dgval
        end
        rhs .*= -1
        # compute adjoint vector
        if OrdinaryDiffEq.isimplicit(alg)
            jac(J, resid, ti, dt, uprev, ui, p)
            λ .= J' \ rhs
        else
            λ .= rhs
        end
        # vector jacobian products
        vjp(uvjpval, pvjpval, λ, ti, dt, uprev, ui, p)
        # accumulate gradient
        dp .-= pvjpval
        # initialize right hand side for the next iteration
        rhs .= uvjpval
    end
    # define gradient with respect to the inputs
    idx = findfirst(tidx -> tidx == t0, t)
    if !isnothing(idx)
        dg(dgval, u0, p, t0, idx)
        du0 .= dgval .- rhs
    end

    return dp, du0
end


text:
include("solver_preamble.jl")
import SeDuMi
factory = with_optimizer(SeDuMi.Optimizer, fid=0)
config = MOI.Test.TestConfig(atol=1e-4, rtol=1e-4, query=false)
@testset "Linear" begin
    Tests.linear_test(factory, config)
end
@testset "SOC" begin
    Tests.soc_test(factory, config, [
        # Expression: termination_status(model) == MOI.INFEASIBLE
        # Evaluated: MathOptInterface.NUMERICAL_ERROR == MathOptInterface.INFEASIBLE
        "sdsos_horn"
    ])
end
@testset "SDP" begin
    Tests.sd_test(factory, config, [
        "quadratic_infeasible_lyapunov_switched_system",
        "quadratic_feasible_lyapunov_switched_system",
        # ALMOST_OPTIMAL
        "maxcut", "chebyshev", "sos_quartic_comparison",
        #   Expression: JuMP.primal_status(model) == MOI.FEASIBLE_POINT
        #    Evaluated: MathOptInterface.NEARLY_FEASIBLE_POINT == MathOptInterface.FEASIBLE_POINT
        "sosdemo5_feasible"
    ])
end


text:
@testset "Pomodoro Stuff Function Tests" begin

  @test isdefined(Pomodoro, :pomodoro_stuff) == true

end

max_stars_repo_name: marcelovmaciel/Pomodoro

max_issues_repo_name: marcelovmaciel/Pomodoro

max_forks_repo_name: marcelovmaciel/Pomodoro


text:
# environment + packages
using Pkg; Pkg.activate(".")
using CSV
using HTTP
using GZip
using GRASS
using LsqFit
using Statistics
using DataFrames
using Interpolations

# plotting
using LaTeXStrings
import PyPlot; plt = PyPlot; mpl = plt.matplotlib; plt.ioff()
using PyCall; animation = pyimport("matplotlib.animation");
mpl.style.use(GRASS.moddir * "figures/fig.mplstyle")

# # define some functions
include(GRASS.moddir * "figures/fig_functions.jl")

# get command line args and output directories
run, plot = parse_args(ARGS)
grassdir, plotdir, datadir = check_plot_dirs()

function download_iag()
    println(">>> Downloading IAG atlas...")
    file = HTTP.download("https://cdsarc.unistra.fr/ftp/J/A+A/587/A65/spvis.dat.gz",
                         GRASS.moddir * "figures/", update_period=Inf)
    println(">>> IAG atlas downloaded!")
    return nothing
end

function read_iag(; isolate=true)
    # download the IAG atlas
    file = GRASS.moddir * "figures/spvis.dat.gz"
    if !isfile(file)
        download_iag()
    end

    # read in the IAG atlas
    iag = GZip.open(file, "r") do io
        CSV.read(io, DataFrame, ignorerepeated=true, delim=" ", header=["wavenum", "nflux", "flux"])
    end

    # convert wavenumber to wavelength in angstroms
    wavs = (1 ./ iag.wavenum) * 1e8

    # reverse to deal with conversion of units
    reverse!(wavs)
    reverse!(iag.nflux)

    # isolate region around 5434.5 line
    if isolate
        ind1 = findfirst(x -> x .> 5435.5, wavs)
        ind2 = findfirst(x -> x .> 5436.55, wavs)
        return wavs[ind1:ind2], iag.nflux[ind1:ind2]
    end
    return wavs, iag.nflux[ind1:ind2]
end

function interpolate_spec(wavs, flux)
    wavs_itp = collect(range(wavs[1], wavs[end], step=mean(diff(wavs))))
    itp = LinearInterpolation(wavs, flux)
    flux_itp = itp.(wavs_itp)
    return wavs_itp, flux_itp
end

# model the iag blends
function model_iag_blends(wavs_sim, flux_sim, wavs_iag, flux_iag; plot=false)
    # roughly align the spectra
    off1 = wavs_iag[argmin(flux_iag)] - wavs_sim[argmin(flux_sim)]
    wavs_iag .-= off1

    # interpolate onto the same grid
    flux_iag ./= maximum(flux_iag)
    itp = LinearInterpolation(wavs_sim, flux_sim, extrapolation_bc=1.0)
    wavs_sim = wavs_iag
    flux_sim = itp.(wavs_sim)

    # calculate some residuals
    resids = flux_iag .- flux_sim

    # models for fit
    @. gaussian(x, a, b, c) = a * exp(-(x - b)^2/(2 * c^2)) + 1
    function tel_model(x, p)
        n = length(p) ÷ 3
        out = ones(length(x))
        for i in 1:n
            out .*= gaussian(x, p[3i-2:3i]...)
        end
        return out .* flux_sim
    end
    # initial guesses for blended lines
    l0 = [-0.005, 5434.0, 0.012]
    l1 = [-0.03, 5434.15, 0.0325]
    l2 = [-0.02, 5434.3, 0.03]
    l3 = [-0.03, 5434.48, 0.02]
    l4 = [-0.02, 5434.55, 0.025]
    l5 = [-0.03, 5434.65, 0.06]
    l6 = [-0.025, 5434.71, 0.02]
    l7 = [-0.05, 5434.812, 0.0305]
    l8 = [-0.035, 5435.01, 0.04]
    l9 = [-0.01, 5435.025, 0.02]
    p0 = [l1..., l2..., l4..., l6..., l7..., l8...]

    # do the fit
    fit = curve_fit(tel_model, wavs_iag, flux_iag, p0)

    # plot diagnostics
    if plot
        fig = plt.figure(figsize=(8,6))
        gs = mpl.gridspec.GridSpec(nrows=2, ncols=1, height_ratios=[2, 1], figure=fig, hspace=0.05)
        ax1 = fig.add_subplot(gs[1])
        ax2 = fig.add_subplot(gs[2])
        ax1.plot(wavs_iag, flux_iag, label="Observed IAG", color="tab:blue")
        ax1.plot(wavs_iag, tel_model(wavs_iag, fit.param), label="Modeled IAG ", ls=":", color="tab:green")
        ax2.scatter(wavs_sim, flux_iag./tel_model(wavs_iag, fit.param), c="k", s=0.5)
        ax1.legend()
        ax1.set_xticklabels([])
        ax1.set_ylabel(L"{\rm Normalized\ Intensity}")
        ax2.set_xlabel(L"{\rm Wavelength\ (\AA)}")
        ax2.set_ylabel(L"{\rm IAG/Model}")
        plt.show()
        plt.clf(); plt.close()
    end
    return flux_iag./(tel_model(wavs_iag, fit.param)./flux_sim)
end

# figure 3 -- compare synthetic and IAG spectra + bisectors
function main()
    # some parameters
    vlim = 2.1e4
    btop = 0.9

    # get spectrum and interpolate onto even wavelength grid
    wavs_iag, flux_iag = read_iag(isolate=true)
    wavs_iag, flux_iag = interpolate_spec(wavs_iag, flux_iag)

    # set up for GRASS spectrum simulation
    lines = [5434.5232]
    depths = [1.0 - minimum(flux_iag)]
    resolution = 700000.0
    spec = SpecParams(lines=lines, depths=depths, resolution=resolution, extrapolate=true)
    disk = DiskParams(N=132, Nt=15)

    # synthesize spectra, calculate ccf, and get CCF bisector
    len = 72
    lambdas1, outspec1 = synthesize_spectra(spec, disk, seed_rng=false, verbose=true, top=NaN)
    outspec1 ./= maximum(outspec1)
    v_grid, ccf1 = calc_ccf(lambdas1, outspec1, spec, normalize=true)
    outspec1 = mean(outspec1, dims=2)[:,1]
    ccfm = mean(ccf1, dims=2)[:,1]
    vel_sim, bis_sim = GRASS.measure_bisector(v_grid, ccfm, interpolate=false, top=btop, len=len)

    # model the line blends out of the IAG spectrum
    println(">>> Modeling out line blends in IAG spectrum...")
    flux_iag_cor = model_iag_blends(lambdas1, outspec1, wavs_iag, flux_iag, plot=false)

    # get offsets to align the spectra in wavelength
    off1 = wavs_iag[argmin(flux_iag)] - lambdas1[argmin(outspec1[:,1]), 1]
    off2 = wavs_iag[argmin(flux_iag_cor)+1] - lambdas1[argmin(outspec1[:,1]), 1]
    wavs_iag .-= off1

    # calculate a CCF for the IAG spectrum and trim it
    v_grid_iag, ccf_iag = calc_ccf(wavs_iag, flux_iag, [wavs_iag[argmin(flux_iag)]],
                                   [1.0 - minimum(flux_iag)],
                                   1e6, normalize=true)
    ind1 = findfirst(x -> x .> -vlim, v_grid_iag)
    ind2 = findfirst(x -> x .> vlim, v_grid_iag)
    vel_iag, bis_iag = GRASS.measure_bisector(v_grid_iag[ind1:ind2], ccf_iag[ind1:ind2],
                                              interpolate=true, top=btop, len=len)

    # calculate a CCF for the cleaned IAG spectrum and trim it
    wavs_iag, flux_iag_cor = interpolate_spec(wavs_iag, flux_iag_cor)
    v_grid_iag2, ccf_iag2 = calc_ccf(wavs_iag, flux_iag_cor, [wavs_iag[argmin(flux_iag_cor)]],
                                    [1.0 - minimum(flux_iag_cor)],
                                    1e6, normalize=true)
    ind1 = findfirst(x -> x .> -vlim, v_grid_iag2)
    ind2 = findfirst(x -> x .> vlim, v_grid_iag2)
    vel_iag2, bis_iag2 = GRASS.measure_bisector(v_grid_iag2[ind1:ind2], ccf_iag2[ind1:ind2],
                                                interpolate=true, top=btop, len=len)

    # interpolate IAG onto same wavelength scale as synthetic spectrum
    itp = LinearInterpolation(wavs_iag, flux_iag, extrapolation_bc=1.0)
    flux_iag_itp = itp.(lambdas1)
    itp = LinearInterpolation(wavs_iag, flux_iag_cor, extrapolation_bc=1.0)
    flux_iag_itp2 = itp.(lambdas1)

    function comparison_plots()
        # overplot the spectra
        fig = plt.figure()
        gs = mpl.gridspec.GridSpec(nrows=2, ncols=1, height_ratios=[2, 1], figure=fig, hspace=0.05)
        ax1 = fig.add_subplot(gs[1])
        ax2 = fig.add_subplot(gs[2])
        ax1.plot(lambdas1, outspec1, c="black", lw= 1.5, label=L"{\rm Synthetic}")
        ax1.plot(wavs_iag, flux_iag./maximum(flux_iag), marker="s", c="tab:blue", ms=2.0, lw=1.0, markevery=10, label=L"{\rm IAG}")
        ax1.plot(wavs_iag, flux_iag_cor./maximum(flux_iag_cor), alpha=0.9, marker="o", c="tab:green", ms=2.0, lw=1.0, markevery=10, label=L"{\rm Cleaned\ IAG}")
        ax2.plot(lambdas1, flux_iag_itp./maximum(flux_iag_itp) .- outspec1, c="tab:blue", marker="s", ms=2.0, lw=0)
        ax2.plot(lambdas1, flux_iag_itp2./maximum(flux_iag_itp2) .- outspec1, c="tab:green", marker="o", ms=2.0, lw=0)

        # set tick labels, axis labels, etc.
        ax1.set_xticklabels([])
        ax1.set_xlim(5434, 5435)
        ax2.set_xlim(5434, 5435)
        ax1.set_ylabel(L"{\rm Normalized\ Intensity}")
        ax2.set_xlabel(L"{\rm Wavelength\ (\AA)}")
        ax2.set_ylabel(L"{\rm IAG\ -\ Synthetic}")
        ax1.legend()
        fig.tight_layout()

        # save the plot
        fig.savefig(plotdir * "fig3a.pdf")
        plt.clf(); plt.close()
        println(">>> Figure written to: " * plotdir * "fig3a.pdf")

        # align bisectors to arbitrary point
        vel_sim .-= mean(vel_sim)
        vel_iag .-= mean(vel_iag) - 42.0
        vel_iag2 .-= mean(vel_iag2)

        # plot the bisectors
        fig = plt.figure()
        gs = mpl.gridspec.GridSpec(nrows=1, ncols=2, width_ratios=[2, 1.1], figure=fig, wspace=0.05)
        ax1 = fig.add_subplot(gs[1])
        ax2 = fig.add_subplot(gs[2])
        ax1.plot(vel_sim, bis_sim, color="black", lw=2.0, label=L"{\rm Synthetic}")
        ax1.plot(vel_iag, bis_iag, marker="s", c="tab:blue", ms=2.0, lw=1.0, label=L"{\rm IAG}")
        ax1.plot(vel_iag2, bis_iag2,marker="o", c="tab:green", ms=2.0, lw=1.0, label=L"{\rm Cleaned\ IAG}")
        ax2.plot(vel_iag .- vel_sim, bis_iag, c="tab:blue", marker="s", ms=2.0, lw=0.0)
        ax2.plot(vel_iag2 .- vel_sim, bis_iag, c="tab:green", marker="o", ms=2.0, lw=0.0)

        # set tick labels, axis labels, etc.
        ax2.set_yticklabels([])
        ax2.yaxis.tick_right()
        ax1.set_ylim(0.1, 1.1)
        ax2.set_xlim(-15, 25)
        ax2.set_ylim(0.1, 1.1)
        ax1.set_xlabel(L"{\rm Relative\ Velocity\ (ms^{-1})}")
        ax1.set_ylabel(L"{\rm Normalized\ Intensity}")
        ax2.set_xlabel(L"{\rm IAG\ -\ Synthetic\ (ms^{-1})}")
        ax1.legend(loc="upper right", prop=Dict("size"=>10), labelspacing=0.25)

        # save the plot
        fig.savefig(plotdir * "fig3b.pdf")
        plt.clf(); plt.close()
        println(">>> Figure written to: " * plotdir * "fig3b.pdf")
        return nothing
    end

    comparison_plots()
    return nothing
end

if (run | plot)
    main()
end

max_stars_repo_name: palumbom/GRASS

max_issues_repo_name: palumbom/GRASS

max_forks_repo_name: palumbom/GRASS


text:
using CSV
using DataDeps
using DataFrames
using Dates
using JLD2

# Obtained link from this webpage:
# https://www.gov.uk/government/statistical-data-sets/price-paid-data-downloads
register(DataDep(
    "uk_house_price_data",
    "This is the UK house price transaction data. Roughly 3.6GB CSV file.",
    "http://prod.publicdata.landregistry.gov.uk.s3-website-eu-west-1.amazonaws.com" *
        "/pp-complete.csv",
))

register(DataDep(
    "uk_postode_lat_lon",
    "UK postcode vs lat-lon mappings.",
    "https://opendata.camden.gov.uk/api/views/77ra-mbbn/rows.csv?accessType=DOWNLOAD",
))

function generate_housing_data(data_dir)

    data_fname = joinpath(data_dir, "housing.jld2")
    if !isfile(data_fname)

        @warn "SERIALISING DATA"

        house_data_file_name = datadep"uk_house_price_data/pp-complete.csv"

        postcode_file_name =
            datadep"uk_postode_lat_lon/National_Statistics_Postcode_Lookup_UK_Coordinates.csv"

        # Load price data.
        raw_data = CSV.read(house_data_file_name, DataFrame; header=false);

        named_data = dropmissing(DataFrame(
            price=raw_data.Column2,
            date=map(x -> parse(Date, x[1:10]), raw_data.Column3),
            post_code=map(x -> x === missing ? x : filter(!isspace, x), raw_data.Column4),
            property_type=raw_data.Column5,
            new_build=raw_data.Column6,
            estate_type=raw_data.Column7,
        ));

        # Load lat lon data.
        raw_lat_lon = CSV.read(postcode_file_name, DataFrame);

        lat_lon = DataFrame(
            post_code=filter.(!isspace, raw_lat_lon[!, "Postcode 1"]),
            lat=raw_lat_lon.Latitude,
            lon=raw_lat_lon.Longitude,
        );

        # Merge the two data sets.
        all_data = innerjoin(named_data, lat_lon; on=:post_code);

        # Serialise the data.
        FileIO.save(data_fname, Dict("data" => all_data))
    end

    # Load the from the serialised format.
    return FileIO.load(data_fname)["data"]
end


text:
#==========================================================
This routine initializes an optimization problem to recover
a CNOT gate on a coupled 3-qubit system. In particular,
    Oscillator A: 2 energy levels, 2 guard states,
    Oscillator B: 2 energy levels, 2 guard states,
    Oscillator S: 1 energy level,  5 guard states,
The drift Hamiltonian in the rotating frame is
    H_0 = - 0.5*ξ_a(a^†a^†aa)
          - 0.5*ξ_b(b^†b^†bb)
          - 0.5*ξ_s(s^†s^†s)
          - ξ_{ab}(a^†ab^†b)
          - ξ_{as}(a^†as^†s)
          - ξ_{bs}(b^†bs^†s).
Here the control Hamiltonian in the rotating frame
includes the usual symmetric and anti-symmetric terms
 H_{sym,1} = p_1(t)(a + a^†), H_{asym,1} = q_1(t)(a - a^†),
 H_{sym,2} = p_2(t)(b + b^†), H_{asym,2} = q_2(t)(b - b^†),
 H_{sym,3} = p_3(t)(s + s^†), H_{asym,3} = q_3(t)(s - s^†),
where a,b,s are the annihilation operators for each qubit.
The problem parameters for this example are,
            ω_a / 2π    = 4.10595      GHz,
            ξ_a / 2π     =  2.198e-02  GHz,
            ω_b / 2π    =  4.81526     GHz,
            ξ_b / 2π     =  2.252e-01  GHz,
            ω_s / 2π    =  7.8447       GHz,
            ξ_s / 2π     =  2.8299e-05 GHz,
            ξ_{ab} / 2π  =  1.0e-06     GHz,
            ξ_{as} / 2π  =  2.494e-03  GHz,
            ξ_{bs} / 2π  =  2.52445e-03 GHz.
We use Bsplines with carrier waves and 3 frequencies per
oscillator:
    Oscillator A: 0, ξ_a, ξ_b
    Oscillator B: 0, ξ_a, ξ_b
    Oscillator S: 0, ξ_{as}, ξ_{bs}.
==========================================================#
using LinearAlgebra
using Ipopt
using Base.Threads
using Random
using DelimitedFiles
using Printf
using FFTW
using Plots
pyplot()
using SparseArrays
using FileIO

#include("Juqbox.jl") # using the consolidated Juqbox module
using Juqbox

Base.show(io::IO, f::Float64) = @printf(io, "%20.13e", f)

function initial_cond(Ntot, N, Ne, Ng)
    Ident = Matrix{Float64}(I, Ntot, Ntot)
    U0 = Ident[1:Ntot,1:N] # initial guess
    #adjust initial guess
    if Ng[1]+Ng[2]+Ng[3] > 0
        Nt = Ne + Ng

        col = 0
        m = 0
        for k3 in 1:Nt[3]
            for k2 in 1:Nt[2]
                for k1 in 1:Nt[1]
                    m += 1
                    # is this a guard level?
                    guard = (k1 > Ne[1]) || (k2 > Ne[2]) || (k3 > Ne[3])
                    if !guard
                        col = col+1
                        U0[:,col] = Ident[:,m]
                    end # if ! guard
                end #for
            end # for
        end # for

    end # if
    return U0
end

fnt = Plots.font("Helvetica", 16)
lfnt = Plots.font("Helvetica", 10)
Plots.default(titlefont=fnt, guidefont=fnt, tickfont=fnt, legendfont=lfnt)

Nosc = 3

Ne1 = 2 # essential energy levels per oscillator # AP: want Ne1=Ne2=2, but Ne3 = 1
Ne2 = 2
Ne3 = 1
Ng1 = 2 # Osc-1, number of guard states
Ng2 = 2 # Osc-2, number of guard states
Ng3 = 5 # Osc-3, number of guard states

Ne = [Ne1, Ne2, Ne3]
Ng = [Ng1, Ng2, Ng3]
Nt = Ne + Ng

N = Ne1*Ne2*Ne3; # Total number of nonpenalized energy levels
Ntot = Nt[1]*Nt[2]*Nt[3]
Nguard = Ntot - N # Total number of guard states

Tmax = 550.0 # 700.0

# frequencies (in GHz, will be multiplied by 2*pi to get angular frequencies in the Hamiltonian matrix)
fa = 4.10595
fb = 4.81526  # official
fs = 7.8447 # storage   # official
rot_freq = [fa, fb, fs] # rotational frequencies
xa = 2 * 0.1099
xb = 2 * 0.1126 # official
xs = 0.002494^2/xa # 2.8298e-5 # official
xab = 1.0e-6 # 1e-6 official
xas = sqrt(xa*xs) # 2.494e-3 # official
xbs = sqrt(xb*xs) # 2.524e-3 # official

# Note: The ket psi = kji> = e_k kron e_j kron e_i.
# We order the elements in the vector psi such that i varies the fastest with i in [1,Nt1], j in [1,Nt2], , k in [1,Nt3]
# The matrix amat = I kron I kron a1 acts on alpha in psi = gamma kron beta kron alpha
# The matrix bmat = I kron a2 kron I acts on beta in psi = gamma kron beta kron alpha
# The matrix cmat = a3 kron I2 kron I1 acts on gamma in psi = gamma kron beta kron alpha

# construct the lowering and raising matricies: amat, bmat, cmat
# and the system Hamiltonian: H0

a1 = Array(Bidiagonal(zeros(Nt[1]),sqrt.(collect(1:Nt[1]-1)),:U))
a2 = Array(Bidiagonal(zeros(Nt[2]),sqrt.(collect(1:Nt[2]-1)),:U))
a3 = Array(Bidiagonal(zeros(Nt[3]),sqrt.(collect(1:Nt[3]-1)),:U))

I1 = Array{Float64, 2}(I, Nt[1], Nt[1])
I2 = Array{Float64, 2}(I, Nt[2], Nt[2])
I3 = Array{Float64, 2}(I, Nt[3], Nt[3])

# create the a, a^\dag, b and b^\dag vectors
amat = kron(I3, kron(I2, a1))
bmat = kron(I3, kron(a2, I1))
cmat = kron(a3, kron(I2, I1))

adag = Array(transpose(amat))
bdag = Array(transpose(bmat))
cdag = Array(transpose(cmat))

# number ops
num1 = Diagonal(collect(0:Nt[1]-1))
num2 = Diagonal(collect(0:Nt[2]-1))
num3 = Diagonal(collect(0:Nt[3]-1))

# number operators
Na = Diagonal(kron(I3, kron(I2, num1)) )
Nb = Diagonal(kron(I3, kron(num2, I1)) )
Nc = Diagonal(kron(num3, kron(I2, I1)) )

H0 = -2*pi*(xa/2*(Na*Na-Na) + xb/2*(Nb*Nb-Nb) + xs/2*(Nc*Nc-Nc) + xab*(Na*Nb) + xas*(Na*Nc) + xbs*(Nb*Nc))

# max coefficient amplitudes, rotating frame
amax = 0.05
bmax = 0.1
cmax = 0.1
maxpar = [amax, bmax, cmax]

# package the lowering and raising matrices together into an one-dimensional array of two-dimensional arrays
# Here we choose dense or sparse representation
use_sparse = true

# dense matrices run faster, but take more memory
Hsym_ops=[Array(amat+adag), Array(bmat+bdag), Array(cmat+cdag)]
Hanti_ops=[Array(amat-adag), Array(bmat-bdag), Array(cmat - cdag)]
H0 = Array(H0)

# Estimate time step
Pmin = 40 # should be 20 or higher
nsteps = calculate_timestep(Tmax, H0, Hsym_ops, Hanti_ops, maxpar, Pmin)

println("Number of time steps = ", nsteps)

samplerate = 32 # for plotting (?)
kpar = 5 # test this component of the gradient

# type of basis function
use_bcarrier = true # false

Ncoupled = length(Hsym_ops)
if use_bcarrier
  Nfreq = 3 # 3 # number of carrier frequencies
else # regular B-splines
  Nfreq=1
  D1 = 400
  nCoeff = 2*Ncoupled*D1 # Total number of B-spline coeffs'. Must be divisible by 6
end

om = zeros(Ncoupled,Nfreq) # In the rotating frame all ctrl Hamiltonians have a zero resonace frequency

# initialize the carrier frequencies
if use_bcarrier
  @assert(Nfreq == 1 || Nfreq == 2 || Nfreq == 3)
  if Nfreq==2
    # freq 2 and 3 for seg 3 and 6 (coupler)
    om[1:Ncoupled,2] .= -1.0*pi*xas # coupling freq for all (re/im)
  elseif Nfreq==3
    # fundamental resonance frequencies for the transmons
      om[1:2,2] .= -2.0*pi*xa # carrier freq's for ctrl Hamiltonian 1 & 2
      om[1:2,3] .= -2.0*pi*xb # carrier freq's for ctrl Hamiltonian 1 & 2
      om[3,2] = -2.0*pi*xas # carrier freq 2 for ctrl Hamiltonian #3
      om[3,3] = -2.0*pi*xbs # carrier freq 2 for ctrl Hamiltonian #3
  end
end
println("Carrier frequencies 1st ctrl Hamiltonian [GHz]: ", om[1,:]./(2*pi))
println("Carrier frequencies 2nd ctrl Hamiltonian [GHz]: ", om[2,:]./(2*pi))
println("Carrier frequencies 3rd ctrl Hamiltonian [GHz]: ", om[3,:]./(2*pi))

casename = "cnot-storage"

# specify target gate
#target for CNOT (2 essential states, 1 g, 1 g, 7 guard states)

# The 2-osc CNOT gate is a 9x4 matrix and the identity for the 3rd oscillator with
# 7 guards is an 9x2 identity matrix

N2tot = Nt[1] * Nt[2]
N2 = Ne[1]*Ne[2]

utarget = zeros(ComplexF64, N2tot*Nt[3], N2*Ne[3])
# target for CNOT gate between oscillators 1 and 2
G2 = zeros(ComplexF64, N2tot, N2)
@assert(Ng[1] == 2 || Ng[1] == 1 || Ng[1] == 0)
@assert(Ne[1] == 2 && Ne[2] == 2)
if Ne[1]==2 && Ne[2] == 2
  if Ng[1] == 0
    G2[1,1] = 1.0
    G2[2,2] = 1.0
    G2[3,4] = 1.0
    G2[4,3] = 1.0
  elseif Ng[1] == 1
    G2[1,1] = 1.0
    G2[2,2] = 1.0
    G2[4,4] = 1.0
    G2[5,3] = 1.0
  elseif Ng[1] == 2
    G2[1,1] = 1.0
    G2[2,2] = 1.0
    G2[5,4] = 1.0
    G2[6,3] = 1.0
  end
end

Ident = Matrix{Float64}(I, Ntot, Ntot)
I3 = Matrix{Float64}(I, Nt[3], Ne[3]);
utarget = kron(I3, G2) # The CNOT is between oscillator 1 and 2. Identity for the 3rd oscillator

# rotation matrices
omega1, omega2, omega3 = Juqbox.setup_rotmatrices(Ne, Ng, rot_freq)

# Compute Ra*Rb*utarget
rot1 = Diagonal(exp.(im*omega1*Tmax))
rot2 = Diagonal(exp.(im*omega2*Tmax))
rot3 = Diagonal(exp.(im*omega3*Tmax))

# target in the rotating frame
vtarget = rot1*rot2*rot3*utarget

U0 = initial_cond(Ntot, N, Ne, Ng)

# NOTE: maxpar is now a vector with 3 elements: amax, bmax, cmax
params = Juqbox.objparams(Ne, Ng, Tmax, nsteps, Uinit=U0, Utarget=vtarget, Cfreq=om, Rfreq=rot_freq,
                          Hconst=H0, Hsym_ops=Hsym_ops, Hanti_ops=Hanti_ops, use_sparse=use_sparse)

Random.seed!(2456)

# setup the initial parameter vector, either randomized or from file
startFromScratch = false # true
startFile="drives/cnot3-pcof-opt.jld2"

if startFromScratch
    D1 = 15 # 20 # number of B-spline coeff per oscillator, freq, p/q
    nCoeff = 2*Ncoupled*Nfreq*D1 # Total number of parameters.
    pcof0 = amax*0.01 * rand(nCoeff)
    println("*** Starting from random pcof with amplitude ", amax*0.01)
else
    #  read initial coefficients from file
    pcof0 = read_pcof(startFile);
    nCoeff = length(pcof0)
    D1 = div(nCoeff, 2*Ncoupled*Nfreq)  # number of B-spline coeff per control function
    nCoeff = 2*Nosc*Nfreq*D1 # just to be safe if the file doesn't contain the right number of elements
    println("*** Starting from B-spline coefficients in file: ", startFile)
end

# min and max B-spline coefficient values
useBarrier = true
# minCoeff, maxCoeff = Juqbox.assign_thresholds(maxpar, Ncoupled, Nfreq, D1)
minCoeff, maxCoeff = Juqbox.assign_thresholds(params,D1,maxpar)

# for ipopt
maxIter = 0 #100 # 0 # 250 #50 # optional argument
lbfgsMax = 250 # optional argument

# output run information
println("*** Settings ***")
println("Frequencies: Alice = ", fa, " Bob = ", fb, " Storage = ", fs)
println("Anharmonic coefficients in the Hamiltonian: xa = ", xa, " xb = ", xb, " xs = ", xs)
println("Coupling coefficients in the Hamiltonian: xab = ", xab, " xas = ", xas, " xbs = ", xbs)
println("Essential states in osc = ", Ne, " Guard states in osc = ", Ng)
println("Total number of states, Ntot = ", Ntot, " Total number of guard states, Nguard = ", Nguard)
println("Number of B-spline parameters per spline = ", D1, " Total number of parameters = ", nCoeff)
println("Max parameter amplitudes: maxpar = ", maxpar)
println("Tikhonov coefficients: tik0 (L2) = ", params.tik0)
if use_sparse
    println("Using a sparse representation of the Hamiltonian matrices")
else
    println("Using a dense representation of the Hamiltonian matrices")
end

wa = Juqbox.Working_Arrays(params,nCoeff)
prob = Juqbox.setup_ipopt_problem(params, wa, nCoeff, minCoeff, maxCoeff, maxIter, lbfgsMax, startFromScratch)

# uncomment to run the gradient checker for the initial pcof
#addOption( prob, "derivative_test", "first-order"); # for testing the gradient

# tmp: test call traceJuqbox()
#objv, objgrad, u_hist, infidelity = Juqbox.traceobjgrad(pcof0, params, true, true);

println("Initial coefficient vector stored in 'pcof0'")


text:
using MDCDL
using Plots

function test_synthesis_speed(ts::TransformSystem, data, N::Integer)
    val, t, bytes, gctime, memallocs = @timed for n = 1:N
        synthesize(ts, data)
    end
    return t
end

# image size
szxs = [ (8,8), (16,16) ,(32,32), (64,64), (128,128), (256,256), (512,512) ]
lnszxs = length(szxs)

nitrs = 100       # Number of elements of iterations

df = (2,2)      # Decimation factor
ord = (2,2)     # Polyphase order
nch = 8         # Number of channels

nsolt = Rnsolt(df, ord, nch)    # create instance of Real-NSOLT (RNSOLT)
rand!(nsolt)    # randomize RNSOLT angle parameteres

tms = Array{Float64}(undef, lnszxs) # elapsed times
for idx = 1:lnszxs
    transform = createTransform(nsolt, Shapes.Vec(szxs[idx]...))
    y = analyze(transform, rand(szxs[idx]...))   # generate dataset
    tms[idx] = test_synthesis_speed(transform, y, nitrs)
end

max_stars_repo_name: nixir/MDCDL

max_issues_repo_name: nixir/MDCDL

max_forks_repo_name: nixir/MDCDL


text:
module JuliaTutor
    using Crayons, Logging, ReplMaker, DataStructures, TerminalMenus
    global repl = nothing
    #Handle lesson plans
    global LESSON_PATH = Base.joinpath( @__DIR__ ,  "Lessons" )
    lessons_available = readdir( LESSON_PATH )
    push!(lessons_available, "Exit Julia Tutor") #provide an easy way out

    #Internal to JuliaTutor
    include("CrayonBox.jl")
    include("CaptureAndEvaluate.jl")
    #External to JuliaTutor IE user facing
    include("Evaluators.jl")
    #export user_experience, Evaluator, ExpressionEvaluator, CommandEvaluator
    include("Tutor.jl")
    #export Lesson, Tutor, display_prompt_and_request
    include("Display.jl")
    #export greet, inform, request
    export help

    global julia_tutor_parser = Tutor("",0,[])
    export julia_tutor_parser

    """
        load_lesson(lesson_location::String)

    reads a structured Lesson file and begins a Julia tutor REPL, parser, etc.
    """
    function load_lesson(lesson_location::String)
        include( lesson_location )

        if !@isdefined(lesson_plan) || (length(lesson_plan) == 0)
            println("The lesson you have attempted to load is lacking a valid `lesson_plan` variable.")
            return nothing
        end

        global julia_tutor_parser = Tutor( lesson_location, 1, lesson_plan )

        function parser_closure(user_input)
            if keywords(julia_tutor_parser, user_input)
                return
            end
            #Display the executed output of the users input
            io = IOBuffer()
            io_context = IOContext( io, :limit => true, :displaysize => (7, 70))
            show(io_context, "text/plain", Meta.eval(Meta.parse(user_input)) );
            println( String( take!( io ) ) )
            #Parse and evaluate the users input
            julia_tutor_parser(user_input)
        end
        #Instantiate a REPL
        global repl = initrepl(
            parser_closure, #ReplMaker doesn't allow functors
            prompt_text="julia tutor> ",
            prompt_color = :yellow,
            start_key=')',
            mode_name="tutor_mode",
            valid_input_checker=complete_julia
        )
        #Enter the REPL without user approval
        ReplMaker.enter_mode!(Base.active_repl.mistate, repl)
        #Display the first lesson prompt
        display_prompt_and_request(julia_tutor_parser)
    end
    export load_lesson

    #Construct main menu.
    """
        menu()

    Greets an end user, but then prompts them to select a lesson plan from /Lessons/
    Once a lesson is selected(by entering a numerical value) it will load that lesson and begin.
    """
    function menu()
        menu_statement = """Please enter the numeral for lesson plan you want to do.
        Note: you can type \"exit()\" and press enter at any time to exit this Julia session and JuliaTutor."""
        println( red_bold(), "> ", white_bold(), menu_statement )
        menut = RadioMenu( lessons_available, pagesize = 7)
        choice = lessons_available[ TerminalMenus.request( "", menut ) ]
        if choice == "Exit Julia Tutor"
            exit()
        else
            load_lesson( Base.joinpath( LESSON_PATH, choice ) )
        end
    end
    export menu

    greet()
end # module