text:
using ChainRulesCore: NoTangent
using Base.Experimental: @opaque

struct ∂⃖rrule{N}; end
struct ∂⃖recurse{N}; end

include("recurse.jl")

function perform_optic_transform(@nospecialize(ff::Type{∂⃖recurse{N}}), @nospecialize(args)) where {N}
    @assert N >= 1

    # Check if we have an rrule for this function
    mthds = Base._methods_by_ftype(Tuple{args...}, -1, typemax(UInt))
    if length(mthds) != 1
        return :(throw(MethodError(ff, args)))
    end
    match = mthds[1]

    mi = Core.Compiler.specialize_method(match)
    ci = Core.Compiler.retrieve_code_info(mi)

    ci′ = copy(ci)
    ci′.edges = MethodInstance[mi]

    r = transform!(ci′, mi.def, length(args) - 1, match.sparams, N)
    if isa(r, Expr)
        return r
    end

    ci′.ssavaluetypes = length(ci′.code)
    ci′.method_for_inference_limit_heuristics = match.method
    ci′
end

# This relies on PartialStruct to infer well
struct Protected{N}
    a
end
(p::Protected)(args...) = getfield(p, :a)(args...)[1]
@Base.aggressive_constprop (::∂⃖{N})(p::Protected{N}, args...) where {N} = getfield(p, :a)(args...)
@Base.aggressive_constprop (::∂⃖{1})(p::Protected{1}, args...) = getfield(p, :a)(args...)
(::∂⃖{N})(p::Protected, args...) where {N} = error("TODO: Can we support this?")

struct OpticBundle{T}
    x::T
    clos
end
Base.getindex(o::OpticBundle, i::Int) = i == 1 ? o.x :
                                        i == 2 ? o.clos :
                                        throw(BoundsError(o, i))
Base.iterate(o::OpticBundle) = (o.x, nothing)
Base.iterate(o::OpticBundle, ::Nothing) = (o.clos, missing)
Base.iterate(o::OpticBundle, ::Missing) = nothing

# Desturucture using `getfield` rather than iterate to make
# inference happier
macro destruct(arg)
    @assert isexpr(arg, :(=))
    lhs = arg.args[1]
    @assert isexpr(lhs, :tuple)
    rhs = arg.args[2]
    s = gensym()
    quote
        $s = $(esc(rhs))
        $([:($(esc(arg)) = getfield($s, $i)) for (i, arg) in enumerate(lhs.args)]...)
        $s
    end
end

@eval function (::∂⃖{1})(::∂⃖{2}, args...)
    @destruct a, b = ∂⃖{3}()(args...)
    (a, $(Expr(:new, Protected{1}, :(Δ->begin
        x1, x2 = b(Δ)
        (x1, $(Expr(:new, Protected{1}, :((Δ...)->begin
            x3, x4 = x2(Δ...)
            (x3, $(Expr(:new, Protected{1}, :(Δ->begin
                x5, x6 = x4(Δ)
                x5, Δ->begin
                    (x7, x8) = x6(Δ...)
                    return (x8, x7)
                end
            end)))), ((x7, x8),)->begin
                (x9, x10) = x8(x7)
                (x10, x9...)
            end
        end)))), ((x9, x10),)->begin
            (x11, x12) = x10(x9...)
            (x12, x11)
        end
    end)))), ((x11, x12),)->begin
        (NoTangent(), x12(x11)...)
    end
end

struct ∂⃖weaveInnerOdd{N, O}; b̄; end
@Base.aggressive_constprop function (w::∂⃖weaveInnerOdd{N, N})(Δ) where {N}
    @destruct c, c̄ = w.b̄(Δ...)
    return (c̄, c)
end
@Base.aggressive_constprop function (w::∂⃖weaveInnerOdd{N, O})(Δ) where {N, O}
    @destruct c, c̄ = w.b̄(Δ...)
    return (c̄, c), ∂⃖weaveInnerEven{plus1(N), O}()
end
struct ∂⃖weaveInnerEven{N, O}; end
@Base.aggressive_constprop function (w::∂⃖weaveInnerEven{N, O})(Δ′, x...) where {N, O}
    @destruct y, ȳ  = Δ′(x...)
    return y, ∂⃖weaveInnerOdd{plus1(N), O}(ȳ)
end

struct ∂⃖weaveOuterOdd{N, O}; end
@Base.aggressive_constprop function (w::∂⃖weaveOuterOdd{N, N})((Δ′′, Δ′′′)) where {N}
    return (NoTangent(), Δ′′′(Δ′′)...)
end
@Base.aggressive_constprop function (w::∂⃖weaveOuterOdd{N, O})((Δ′′, Δ′′′)) where {N, O}
    @destruct α, ᾱ = Δ′′′(Δ′′)
    return (NoTangent(), α...), ∂⃖weaveOuterEven{plus1(N), O}(ᾱ)
end
struct ∂⃖weaveOuterEven{N, O}; ᾱ end
@Base.aggressive_constprop function (w::∂⃖weaveOuterEven{N, O})(Δ⁴...) where {N, O}
    return w.ᾱ(Base.tail(Δ⁴)...), ∂⃖weaveOuterOdd{plus1(N), O}()
end

function (::∂⃖{N})(::∂⃖{1}, args...) where {N}
    @destruct (a, ā) = ∂⃖{plus1(N)}()(args...)
    let O = c_order(N)
        (a, Protected{N}(@opaque Δ->begin
                (b, b̄) = ā(Δ)
                b, ∂⃖weaveInnerOdd{1, O}(b̄)
            end
        )), ∂⃖weaveOuterOdd{1, O}()
    end
end

function (::∂⃖{N})(::∂⃖{M}, args...) where {N, M}
    # TODO
    @destruct (a, b) = ∂⃖{N+M}()(args...)
    error("Not implemented yet ($N, $M)")
end

macro OpticBundle(a, b)
    aa = gensym()
    esc(quote
        $aa = $a
        $(Expr(:new, :(OpticBundle{typeof($aa)}), aa, b))
    end)
end

# ∂⃖rrule has a 4-recurrence - we model this as 4 separate structs that we
# cycle between. N.B.: These names match the names that these variables
# have in Snippet 19 of the terminology guide. They are probably not ideal,
# but if you rename them here, please update the terminology guide also.

struct ∂⃖rruleA{N, O}; ∂; ȳ; ȳ̄ ; end
struct ∂⃖rruleB{N, O}; ᾱ; ȳ̄ ; end
struct ∂⃖rruleC{N, O}; ȳ̄ ; Δ′′′; β̄ ; end
struct ∂⃖rruleD{N, O}; γ̄; β̄ ; end

@Base.aggressive_constprop function (a::∂⃖rruleA{N, O})(Δ) where {N, O}
    # TODO: Is this unthunk in the right place
    @destruct (α, ᾱ) = a.∂(a.ȳ, unthunk(Δ))
    (α, ∂⃖rruleB{N, O}(ᾱ, a.ȳ̄))
end

@Base.aggressive_constprop function (b::∂⃖rruleB{N, O})(Δ′...) where {N, O}
    @destruct ((Δ′′′, β), β̄) = b.ᾱ(Δ′)
    (β, ∂⃖rruleC{N, O}(b.ȳ̄, Δ′′′, β̄))
end

@Base.aggressive_constprop function (c::∂⃖rruleC{N, O})(Δ′′) where {N, O}
    @destruct (γ, γ̄) = c.ȳ̄((Δ′′, c.Δ′′′))
    (Base.tail(γ), ∂⃖rruleD{N, O}(γ̄, c.β̄))
end

@Base.aggressive_constprop function (d::∂⃖rruleD{N, O})(Δ⁴...) where {N, O}
    (δ₁, δ₂), δ̄  = d.γ̄(ZeroTangent(), Δ⁴...)
    (δ₁, ∂⃖rruleA{N, O+1}(d.β̄ , δ₂, δ̄ ))
end

# Terminal cases
@Base.aggressive_constprop function (c::∂⃖rruleB{N, N})(Δ′...) where {N}
    @destruct (Δ′′′, β) = c.ᾱ(Δ′)
    (β, ∂⃖rruleC{N, N}(c.ȳ̄, Δ′′′, nothing))
end
@Base.aggressive_constprop (c::∂⃖rruleC{N, N})(Δ′′) where {N} =
    Base.tail(c.ȳ̄((Δ′′, c.Δ′′′)))
(::∂⃖rruleD{N, N})(Δ...) where {N} = error("Should not be reached")

# ∂⃖rrule
@Base.pure term_depth(N) = 2^(N-2)
function (::∂⃖rrule{N})(z, z̄) where {N}
    @destruct (y, ȳ) = z
    y, ∂⃖rruleA{term_depth(N), 1}(∂⃖{minus1(N)}(), ȳ, z̄)
end

function (::∂⃖{N})(f::Core.IntrinsicFunction, args...) where {N}
    # A few intrinsic functions are inserted by the compiler, so they need to
    # be handled here. Otherwise, we just throw an appropriate error.
    if f === Core.Intrinsics.not_int && length(args) == 1
        return f(args...), EvenOddOdd{1, c_order(N)}(
            Δ->(NoTangent(), NoTangent()),
            Δ->NoTangent())
    end

    error("Rewrite reached intrinsic function $f. Missing rule?")
end

# The static parameter on `f` disables the compileable_sig heuristic
function (::∂⃖{N})(f::T, args...) where {T, N}
    if N == 1
        # Base case (inlined to avoid ambiguities with manually specified
        # higher order rules)
        z = rrule(f, args...)
        if z === nothing
            return ∂⃖recurse{1}()(f, args...)
        end
        return z
    else
        ∂⃖p = ∂⃖{minus1(N)}()
        @destruct z, z̄ = ∂⃖p(rrule, f, args...)
        if z === nothing
            return ∂⃖recurse{N}()(f, args...)
        else
            return ∂⃖rrule{N}()(z, z̄)
        end
    end
end

@Base.pure function (::∂⃖{1})(::typeof(Core.apply_type), head, args...)
    return rrule(Core.apply_type, head, args...)
end

struct KwFunc{T,S}
    f::T
    kwf::S
    function KwFunc(f)
        kwf = Core.kwfunc(f)
        new{Core.Typeof(f), Core.Typeof(kwf)}(f, kwf)
    end
end
(kw::KwFunc)(args...) = kw.kwf(args...)

function ChainRulesCore.rrule(::typeof(Core.kwfunc), f)
    KwFunc(f), Δ->(NoTangent(), Δ)
end

function ChainRulesCore.rrule(::KwFunc, kwargs, f, args...)
    r = Core.kwfunc(rrule)(kwargs, rrule, f, args...)
    if r === nothing
        return nothing
    end
    x, back = r
    x, Δ->begin
        (NoTangent(), NoTangent(), back(Δ)...)
    end
end

@Base.aggressive_constprop function ChainRulesCore.rrule(::typeof(Core.getfield), s, field::Symbol)
    getfield(s, field), let P = typeof(s)
        @Base.aggressive_constprop Δ->begin
            nt = NamedTuple{(field,)}((Δ,))
            (NoTangent(), Tangent{P, typeof(nt)}(nt), NoTangent())
        end
    end
end

struct ∂⃖getfield{n, f}; end
@Base.aggressive_constprop function (::∂⃖getfield{n, f})(Δ) where {n,f}
    if @generated
        return Expr(:call, tuple, NoTangent(),
            Expr(:call, tuple, (i == f ? :(Δ) : ZeroTangent() for i = 1:n)...),
            NoTangent())
    else
        return (NoTangent(), ntuple(i->i == f ? Δ : ZeroTangent(), n), NoTangent())
    end
end

struct EvenOddEven{O, P, F, G}; f::F; g::G; end
EvenOddEven{O, P}(f::F, g::G) where {O, P, F, G} = EvenOddEven{O, P, F, G}(f, g)
struct EvenOddOdd{O, P, F, G}; f::F; g::G; end
EvenOddOdd{O, P}(f::F, g::G) where {O, P, F, G} = EvenOddOdd{O, P, F, G}(f, g)
@Base.aggressive_constprop (o::EvenOddOdd{O, P, F, G})(Δ) where {O, P, F, G} = (o.f(Δ), EvenOddEven{plus1(O), P, F, G}(o.f, o.g))
@Base.aggressive_constprop (e::EvenOddEven{O, P, F, G})(Δ...) where {O, P, F, G} = (e.g(Δ...), EvenOddOdd{plus1(O), P, F, G}(e.f, e.g))
@Base.aggressive_constprop (o::EvenOddOdd{O, O})(Δ) where {O} = o.f(Δ)

@Base.aggressive_constprop function (::∂⃖{N})(::typeof(Core.getfield), s, field::Int) where {N}
    getfield(s, field), EvenOddOdd{1, c_order(N)}(
        ∂⃖getfield{nfields(s), field}(),
        @Base.aggressive_constprop (_, Δ, _)->getfield(Δ, field))
end

@Base.aggressive_constprop function (::∂⃖{N})(::typeof(Base.getindex), s::Tuple, field::Int) where {N}
    getfield(s, field), EvenOddOdd{1, c_order(N)}(
        ∂⃖getfield{nfields(s), field}(),
        @Base.aggressive_constprop (_, Δ, _)->lifted_getfield(Δ, field))
end

function (::∂⃖{N})(::typeof(Core.getfield), s, field::Symbol) where {N}
    getfield(s, field), let P = typeof(s)
        EvenOddOdd{1, c_order(N)}(
            (@Base.aggressive_constprop Δ->begin
                nt = NamedTuple{(field,)}((Δ,))
                (NoTangent(), Tangent{P, typeof(nt)}(nt), NoTangent())
            end),
            (@Base.aggressive_constprop (_, Δs, _)->begin
                isa(Δs, Union{ZeroTangent, NoTangent}) ? Δs : getfield(ChainRulesCore.backing(Δs), field)
            end))
    end
end

# TODO: Temporary - make better
function (::∂⃖{N})(::typeof(Base.getindex), a::Array, inds...) where {N}
    getindex(a, inds...), let
        EvenOddOdd{1, c_order(N)}(
            (@Base.aggressive_constprop Δ->begin
                Δ isa AbstractZero && return (NoTangent(), Δ, map(Returns(Δ), inds)...)
                BB = zero(a)
                BB[inds...] = Δ
                (NoTangent(), BB, map(x->NoTangent(), inds)...)
            end),
            (@Base.aggressive_constprop (_, Δ, _)->begin
                getindex(Δ, inds...)
            end))
    end
end

struct tuple_back{M}; end
(::tuple_back)(Δ::Tuple) = Core.tuple(NoTangent(), Δ...)
(::tuple_back{N})(Δ::AbstractZero) where {N} = Core.tuple(NoTangent(), ntuple(i->Δ, N)...)
(::tuple_back{N})(Δ::Tangent) where {N} = Core.tuple(NoTangent(), ntuple(i->lifted_getfield(Δ, i), N)...)

function (::∂⃖{N})(::typeof(Core.tuple), args::Vararg{Any, M}) where {N, M}
    Core.tuple(args...),
        EvenOddOdd{1, c_order(N)}(
            tuple_back{M}(),
            (Δ...)->Core.tuple(Δ[2:end]...)
        )
end

struct UnApply{Spec, Types}; end
@generated function (::UnApply{Spec, Types})(Δ) where {Spec, Types}
    args = Any[NoTangent(), NoTangent(), :(Δ[1])]
    start = 2
    for (l, T) in zip(Spec, Types.parameters)
        if T <: Array
            push!(args, :([Δ[$(start:(start+l-1))]...]))
        else
            push!(args, :(Δ[$(start:(start+l-1))]))
        end
        start += l
    end
    :(Core.tuple($(args...)))
end

struct ApplyOdd{O, P}; u; ∂⃖f; end
struct ApplyEven{O, P}; u; ∂⃖∂⃖f; end
@Base.aggressive_constprop function (a::ApplyOdd{O, P})(Δ) where {O, P}
    r, ∂⃖∂⃖f = a.∂⃖f(Δ)
    (a.u(r), ApplyEven{plus1(O), P}(a.u, ∂⃖∂⃖f))
end
@Base.aggressive_constprop function (a::ApplyEven{O, P})(_, _, ff, args...) where {O, P}
    r, ∂⃖∂⃖∂⃖f = Core._apply_iterate(iterate, a.∂⃖∂⃖f, (ff,), args...)
    (r, ApplyOdd{plus1(O), P}(a.u, ∂⃖∂⃖∂⃖f))
end
@Base.aggressive_constprop function (a::ApplyOdd{O, O})(Δ) where {O}
    r = a.∂⃖f(Δ)
    a.u(r)
end

function (this::∂⃖{N})(::typeof(Core._apply_iterate), iterate, f, args::Union{Tuple, Vector, NamedTuple}...) where {N}
    @assert iterate === Base.iterate
    x, ∂⃖f = Core._apply_iterate(iterate, this, (f,), args...)
    return x, ApplyOdd{1, c_order(N)}(UnApply{map(length, args), typeof(args)}(), ∂⃖f)
end

@Base.pure c_order(N::Int) = 2^N - 1

@Base.pure function (::∂⃖{N})(::typeof(Core.apply_type), head, args...) where {N}
    Core.apply_type(head, args...), NonDiffOdd{plus1(plus1(length(args))), 1, c_order(N)}()
end

@Base.aggressive_constprop lifted_getfield(x, s) = getfield(x, s)
lifted_getfield(x::ZeroTangent, s) = ZeroTangent()
lifted_getfield(x::NoTangent, s) = NoTangent()

function lifted_getfield(x::Tangent, s)
    z = getfield(ChainRulesCore.backing(ChainRulesCore.canonicalize(x)), s)
    z
end

function lifted_getfield(x::Tangent{<:Tangent{T}}, s) where T
    bb = getfield(x.backing, 1)
    z = lifted_getfield(bb, s)
    z
end

ChainRulesCore.backing(::ZeroTangent) = ZeroTangent()
ChainRulesCore.backing(::NoTangent) = NoTangent()

function reload()
    Core.eval(Diffractor, quote
        function (ff::∂⃖recurse)(args...)
            $(Expr(:meta, :generated_only))
            $(Expr(:meta,
                    :generated,
                    Expr(:new,
                        Core.GeneratedFunctionStub,
                        :perform_optic_transform,
                        Any[:ff, :args],
                        Any[],
                        @__LINE__,
                        QuoteNode(Symbol(@__FILE__)),
                        true)))
        end
    end)
end
reload()


text:
export compute_dt

"""
    compute_dt( sqg, ogird, w)

CFL of the (large-scale) advection

"""
function compute_dt(sqg)

    update_velocities!(sqg)
    dx = sqg.grid.dx
    dy = sqg.grid.dy

    bound1 = maximum(abs.(sqg.u_x .* π ./ dx) .+ abs.(sqg.u_y .* π ./ dy))
    bound1 = 1 / bound1 / 4

    # CFL of the hyperviscosity

    dx2 = ((dx, dy) ./ pi) .^ 2
    bound2 = 1 / sqg.hv_val * (prod(dx2) / sum(dx2))^(sqg.hv_order ÷ 2)

    # Minimum of the CFL
    min(minimum(bound1), bound2)

end

ext: jl


text:
# module Preconditioners

# using JuLIP, LinearAlgebra

# @doc raw"""
# `AdjacancyPrecon` :

# WARNING: While the code accounts for PBC, the following description does not.

# If `(r0, r1, k)` is in the `intervals` list and if
# `r0 <= rij <= r1` then the 3 x 3 block
# ```math
#    P_{ij} = - k \big( (1-c) \hat{R}_{ij} \otimes \hat{R}_{ij} + c I \big),
# ```
# where ``\hat{R}_{ij}`` is the bond direction and `c = innerstab`. The diagonal
# is then obtained by
# ```math
#    P_{ii} = \sum_{j \neq i} P_{ij} + c_g I
# ```
# where `cg = stab` is a second, global, stabilisation constant.

# Regarding the choice of `innerstab`:
# * Choose `innerstab = 0` to get full infinitesimal rotation-invariance
# * Choose `innerstab = 1` to get full resistance against infinitesimal rotations
# * In optimisation we often found that `innerstab = 0.1` is a decent choice.
# """
# struct AdjacancyPrecon
#    intervals::Vector{NamedTuple{(:r0, :r1, :k), Tuple{Float64, Float64, Float64}}}
#    stab::Float64
#    innerstab::Float64
# end

# import JuLIP.FIO: read_dict, write_dict

# AdjacancyPrecon(intervals; stab = 1e-4, innerstab = 0.1) =
#    AdjacancyPrecon(intervals, stab, innerstab)

# _cutoff(P::AdjacancyPrecon) =
#    maximum( i.r1 for i in P.intervals )

# function (P::AdjacancyPrecon)(R::JVecF)
#    r = norm(R)
#    R̂ = R / r
#    Π = (1 - P.innerstab) * R̂ * R̂' + P.innerstab * one(JMat{Float64})
#    for i in P.intervals
#       if i.r0 <= r <= i.r1
#          return i.k * Π
#       end
#    end
#    return 0.0 * Π
# end

# function (precon::AdjacancyPrecon)(at::Atoms)
#    P = zeros(3*length(at), 3*length(at))
#    nlist = neighbourlist(at, _cutoff(precon))
#    for i = 1:length(at)
#       Js, Rs = neigs(nlist, i)
#       for (j, Rij) in zip(Js, Rs)
#          if j < i; continue; end
#          Pij = precon(Rij)
#          inds = 3*(i-1) .+ (1:3)
#          jnds = 3*(j-1) .+ (1:3)
#          for a = 1:3, b = 1:3
#             P[inds[a], jnds[b]] -= Pij[a,b]
#             P[jnds[b], inds[a]] -= Pij[a,b]
#             P[inds[a], inds[b]] += Pij[a,b]
#             P[jnds[a], jnds[b]] += Pij[a,b]
#          end
#       end
#    end
#    P += precon.stab * I
#    return P
# end

# end

#---

#---
# CAS - please convert this into a test set:

@info("Test evaluation and some basic properties of AdjacancyPrecon")

using JuLIP, Test

at = bulk(:W, cubic=true) * 2

intervals = [ (r0 = 2.0, r1 = 3.5, k = 1.0),
              (r0 = 3.5, r1 = 4.5, k = 0.1) ]

# the block-Identity preconditoner
adjp1 =  Preconditioners.AdjacancyPrecon(intervals; innerstab = 1.0)
# a variation that accounts for some infinitesimal rotational invariance
# the parameter innerstab can interpolate between the two.
adjp2 =  Preconditioners.AdjacancyPrecon(intervals; innerstab = 0.1)  # <= DEFAULT
adjp3 =  Preconditioners.AdjacancyPrecon(intervals; innerstab = 0.0)

for adjp in [adjp1, adjp2, adjp3]
   P = adjp(at)
   println(@test P' ≈ P)
   sumP = sum(P, dims=1)
   println(@test all(sumP .≈ adjp.stab))
   sumP = sum(P, dims=2)
   println(@test norm(sumP .- adjp.stab, Inf) < 1e-12)
   spec = sort(eigvals(P))
   println(@test all(spec[1:3] .≈ adjp.stab))
end

#---
# # how to use:
# weights = Dict()
# weights["precon"] = Dict("default" => adjp)


text:
"""
    surrogate_model(plan::AbstractArray{T,2}, samples::AbstractArray{T,2}; options=Options()) where T

Returns a surrogate model function based on an optimized Radial Basis Function
interpolant. Depending on the supplied options; the kernel, kernel width and scaling of
input data is optimized.

...
# Arguments
- `plan::AbstractArray{T,2}`:
    sample locations where each column corresponds to the location of one point.
    `size(plan) = (num_dimensions,num_samples)`.
- `samples::AbstractArray{T,2}`:
    function value at each sample location. each column contains one value from
    the corresponding plan location. `size(samples) = (1,num_samples)`.
- `options=Options()`:
    all options available to customize the surrogate optimization.
...
"""
function surrogate_model(plan::AbstractArray{T,2}, samples::AbstractArray{T,2}; options::Options=Options()) where T

    @unpack num_interpolants, trace, parallel_surrogate, smooth = options

    (trace == :verbose) && println("Creating optimized surrogate model")

    (length(samples) != size(plan,2)) && error("plan and samples do not have the correct sizes")

    @timeit_debug "hypers opt" begin
    #Optimize RBF hypers for the ensamble of interpolants
    if parallel_surrogate
        mres = pmap(  (x)->rbf_hypers_opt(samples, plan, options),
                        1:num_interpolants)
    else
        mres = map(  (x)->rbf_hypers_opt(samples, plan, options),
                        1:num_interpolants)
    end
    end

    optres = first.(mres)
    fitness = last.(mres)

    #Print hyperparameter results
    (trace == :verbose) && print_sm_opt(fitness,optres,smooth)

    # Sample order for ScatteredInterpolation
    samples_vec = vec(samples)

    # Interpolation object based on the optimization results
    itp = Array{ScatteredInterpolation.RBFInterpolant,1}(undef,num_interpolants)
    for i = 1:num_interpolants
        # Preprocess the points based on the settings used.
        preprocessed_point = preprocess_point(plan,optres[i].scaling,base_scale=plan)

        itp[i] = interpolate(optres[i].kernelFunc, preprocessed_point,
                            samples_vec, smooth=optres[i].smooth)
    end

    # Pre allocate and calculate to reduce cost of evaluating surrogate model.
    preprocessed_est_point = Array{Float64,2}(undef,size(plan,1),1)
    old_min = minimum(plan,dims=2)
    old_max = maximum(plan,dims=2)

    sm_func = (estimation_point) ->  surrogate_evaluate.(Ref(preprocessed_est_point),
                                        Ref(estimation_point),Tuple(itp),Tuple(optres),
                                        Ref(old_min),Ref(old_max))

    return sm_func, optres
end

function print_sm_opt(fitness,optres,smooth)
    print(@sprintf("    mean surrogate fitness (smaller is better)= %.5g,  (min = %.5g, max = %.5g)\n",mean(fitness),minimum(fitness),maximum(fitness)))

    sc_res = [opt.scaling./opt.scaling[1] for opt in optres]
    println("    mean dim. scale (norm. to first dim.) = ", mean(sc_res),
                                        " (min = ",minimum(sc_res),
                                        ", max = ",maximum(sc_res),")")
    sm_res = [opt.smooth for opt in optres]
    if smooth == :variable
        println("    mean smoothing = ", mean.(sm_res),
                                        " (min = ",minimum.(sm_res),
                                        ", max = ",maximum.(sm_res),") \n")
    else
        println(@sprintf("    mean smoothing = %.5g,  (min = %.5g, max = %.5g)\n",mean(sm_res),minimum(sm_res),maximum(sm_res)))
    end
end


text:
using DrWatson
@quickactivate "ExercisesRepo"
include(srcdir("style.jl"))
using DynamicalSystems, PyPlot, Random

# %% 1d climate
using DynamicalSystems, PyPlot, Roots

αtan(T) = 0.5 - (0.4/π)*atan((T-263)/2)
dTdt(T, ε = 0.65, α=αtan, s= 1.0) = s*(1 - α(T)) - 1.6e-10 * ε * T^4
dTdt(T; ε = 0.65, α=αtan, s = 1.0) = dTdt(T, ε, α, s)

# Ts = 200:400.0
# plot(Ts, dTdt.(Ts))
# plot(Ts, dTdt.(Ts, 0.2))
# plot(Ts, dTdt.(Ts, 0.9))
# axhline(0)

fig = figure(figsize = (figx/1.5, figx/3))
Ts = 200:0.5:320.0
arrows = 210:10:300 |> collect
deleteat!(arrows, findfirst(isequal(260), arrows))
roots = Roots.find_zeros(dTdt, Ts[1], Ts[end])
plot(Ts, dTdt.(Ts), color = "C1", label = "\$dT/dt\$")
axhline(0; lw = 2.0, zorder = -99)
xlim(Ts[1], Ts[end])
ylim(-0.2, 0.2)
yticks([-0.1, 0, 0.1])
xlabel("\$T\$")
tight_layout()
for (i, r) in enumerate(roots)
    plot(r, 0, marker = "o", markeredgecolor = "k", markerfacecolor = iseven(i) ? "w" : "k",
    markersize = 15, mew = 2)
end
for r in arrows
    f = dTdt(r)
    x, dx = f > 0 ? (r - 5, 5) : (r+5, -5)
    ff = abs(1.2f)^2 + 0.01
    arrow(x, 0, dx, 0; color = "C2", width = ff, length_includes_head = false,
    head_width = 1.5ff, head_length = 100ff
    )
end
legend()
tight_layout()
subplots_adjust(bottom = 0.2, left = 0.1)
wsave(plotsdir("1dstatespace"), fig)

# %% Power spectra for three orbits of Henon-Heiles
using DynamicalSystems, PyPlot, FFTW, Statistics

u0s = (
    [0.0, -0.25, 0.42, 0.0],
    [0.0, 0.30266571044921875, 0.4205654433900762, 0.0],
    [0.0, 0.1, 0.5, 0.0],
)

labels = (
    "chaotic",
    "periodic",
    "quasiperiodic",
)

hh = Systems.henonheiles()
fig, axs = subplots(3, 1; figsize = (figx/2, 2figy), sharex = true)

δt = 0.05

for (i, u) in enumerate(u0s)
   # r = trajectory(hh, 1000.0, u; dt = 0.1)[:, 1]
    r = trajectory(hh, 30000.0, u; dt = δt)[:, 1]
    P = abs2.(rfft(r .- mean(r)))
    P[1] = P[2]
    ν = rfftfreq(length(r))/δt
    # axs[i].plot(ν, P ./ maximum(P),
    axs[i].semilogy(ν, P ./ maximum(P),
    label = labels[i], linewidth = i == 1 ? 1.0 : 2.0, color = "C$(i-1)")
    axs[i].text(0.99, 0.8, labels[i]; ha = "right", transform = axs[i].transAxes,
    color = "C$(i-1)")
    # @show std(r)
    # r .+= 0.5randn(length(r))
    # P = abs2.(rfft(r .- mean(r)))
    # PyPlot.plot(10rfftfreq(length(r)), P ./ maximum(P),
    # lw = 1.0, alpha = 0.5, color = "C$(i-1)")
    # PyPlot.plot(r)
    # axs[i].set_yticks([])
    axs[i].set_ylim(10.0^(-4.5), 1.0)
    axs[i].set_yticks(10.0 .^ (0:-2:-5))
    axs[i].set_xlim(0, 0.4)
end
axs[2].set_ylabel("\$P / \\mathrm{max}(P)\$")
axs[3].set_xlabel("frequency \$\\nu\$")
fig.tight_layout()
fig.subplots_adjust(bottom = 0.1, left = 0.2, top = 0.98, right = 0.93, hspace = 0.1)
wsave(plotsdir("spectra"), fig)

# %% fitzhugh new
fig, axs = subplots(2,3; figsize = (figx, 2figy))
fig.tight_layout()

a = 3.
b = 0.2
ε = 0.01
I = 0.0

ds = Systems.fitzhugh_nagumo([0,0.0]; a, b, ε, I)

xgrid = -0.3:0.02:1.1
ygrid = -0.1:0.01:0.4

function add_nullclines!(ax, a, b, ε, I)
    u = -0.3:0.01:1.1
    w1 = @. u
    w2 = @. a*u*(u-b)*(1-u) + I
    ax.plot(u, w1, "--", color = "C3")
    ax.plot(u, w2, ":", color = "C4")
    ax.set_ylim(ygrid[1], ygrid[end])    # why is this necessary?
end

function add_streamlines!(ax, a, b, ε, I)
    ux = zeros(length(xgrid), length(ygrid))
    uy = copy(ux)
    set_parameter!(ds, [a, b, ε, I])
    for (i, x) in enumerate(xgrid)
        for (j, y) in enumerate(ygrid)
            ux[i, j], uy[i, j] = ds.f(SVector(x, y), ds.p, 0)
        end
    end
    ax.streamplot(Vector(xgrid), Vector(ygrid), ux', uy'; linewidth = 0.5, density = 0.5, color = "C7")
    ax.set_ylim(ygrid[1], ygrid[end])
end

# Fig A (top left)
# ================
ax = axs[1]
add_nullclines!(ax, a, b, ε, I)
add_streamlines!(ax, a, b, ε, I)

# Add tex to axs [1] explaining nullclines
sss = 24
props = Dict(:boxstyle=>"round", :alpha=>1., :facecolor => "white")
ax.text(0.5, 0.35, "\$\\dot{w}>0\$", color = "C3", size = sss, bbox=props)
ax.text(-0.1, 0.35, "\$\\dot{w}<0\$", color = "C3", size = sss, bbox=props)
ax.text(0, 0.1, "\$\\dot{u}<0\$", color = "C4", size = sss, bbox=props)
ax.text(0.45, -0.05, "\$\\dot{u}>0\$", color = "C4", size = sss, bbox=props)
ax.set_ylabel("\$w\$"; labelpad = -15)
ax.set_xlabel("\$u\$"; labelpad = -15)
ax.set_xticks([0, 1])
ax.set_yticks([0, 0.4])

# Fig B & C: two trajectries & zoomin
# ==================
for ax in (axs[3], axs[5])
    add_nullclines!(ax, a, b, ε, I)

    u1 = [0.19, 0.]
    tr1 = trajectory(ds, 200.0, u1)
    ax.plot(columns(tr1)...; color = "C2", lw = 1.5)
    ax.scatter(tr1[1]...; color = "C2", s = 20)

    u2 = [0.21, 0.]
    tr2 = trajectory(ds, 300.0, u2)
    ax.plot(columns(tr2)...; color = "C1", lw = 1.5)
    ax.scatter(tr2[1]...; color = "C1", s = 20)

    # ic = (tr1[1] + tr2[1]) ./ 2
    # ax.scatter(ic...; color = "k", s = 10)

    ax.plot(0, 0; marker = "o", mec = "C0", mew = 1,
        markersize = 6, mfc = "w", zorder = 99)
    ax.axhline(0; color = "k", lw = 0.5, zorder = 1)
end
ax = axs[3]
ax.set_ylabel("\$w\$"; labelpad = -15)
ax.set_xlabel("\$u\$"; labelpad = -15)
ax.set_xticks([0, 1])
ax.set_yticks([0, 0.4])

zbox = ((-0.05, -0.04), (0.25, 0.04))
axis_zoomin!(axs[5],axs[3],  zbox, zbox, "C0"; lw = 1.0)
axs[5].set_xlim(zbox[1][1], zbox[2][1])
axs[5].set_ylim(zbox[1][2], zbox[2][2])
axs[5].set_xticks([])
axs[5].set_yticks([])

# Fig D: timeseries, pulses, refractory period
using OrdinaryDiffEq
pulses_start = [20, 80, 170] # callback times
pulse_width = 4
pulses_end = pulses_start .+ pulse_width
pulses = sort!(vcat(pulses_start, pulses_end))
Ipulse = 0.2
condition(u,t,integ) = t ∈ pulses # pulse times
function affect!(integ)
    i = integ.t ∈ pulses_start ? Ipulse : 0.0
    integ.p[4] = i
end
cb = DiscreteCallback(condition, affect!)

Tf = 250.0
dt = 0.1
prob = ODEProblem(ds, (0.0, Tf))
sol = solve(prob, Tsit5(); callback=cb, tstops = pulses, dtmax = 0.1)

axs[2].plot(sol.t, sol[1, :])
pt = [any(x -> x ≤ t ≤ x + pulse_width, pulses_start) ? Ipulse : 0.0 for t in sol.t]
axs[2].plot(sol.t, pt; color = "C2")
axs[2].set_xlabel("\$t\$"; labelpad = -15)
axs[2].set_ylabel("\$u\$"; labelpad = -15)
axs[2].set_yticks([0, 1])
axs[2].set_xlim(0, Tf)
axs[2].set_xticks([pulses_start..., Tf])

# Fig E: multi stable
a = 8.
b = 0.2
ε = 0.01
I = 0.0

xgrid = -0.3:0.02:1.1
ygrid = -0.2:0.01:1

add_nullclines!(axs[4], a, b, ε, I)
set_parameter!(ds, [a,b,ε,I])
ax = axs[4]

u1 = [0.3, 0.]
tr1 = trajectory(ds, 300.0, u1)
ax.plot(columns(tr1)...; color = "C2")
ax.scatter(tr1[1]...; color = "C2", s = 20)

u2 = [0.55, 0.8]
tr2 = trajectory(ds, 300.0, u2)
ax.plot(columns(tr2)...; color = "C1")
ax.scatter(tr2[1]...; color = "C1", s = 20)

ax.set_ylabel("\$w\$"; labelpad = -15)
ax.set_xlabel("\$u\$"; labelpad = -15)
ax.set_xticks([0, 1])
ax.set_yticks([0, 1])

# Fig F: limit cycle
# -------------------------
ax = axs[6]
a = 3.
b = -0.05
ε = 0.01
I = 0.
set_parameter!(ds, [a,b,ε,I])
add_nullclines!(ax, a, b, ε, I)

tr = trajectory(ds, 400.0, [0.0, 0.1])
ax.plot(columns(tr)...; color = "C2")
ax.scatter(tr[1]...; color = "C2", s = 20)
ax.set_xlim(-0.5,1.1)
ax.set_ylim(-0.1,0.6)
ax.set_ylabel("\$w\$"; labelpad = -15)
ax.set_xlabel("\$u\$"; labelpad = -15)
ax.set_xticks([0, 1])
ax.set_yticks([0, 0.6])

fig.subplots_adjust(bottom = 0.08, left = 0.05, top = 0.95, right = 0.97, hspace = 0.2)
add_identifiers!(fig)
# wsave(plotsdir("fitzhugh"), fig)

max_stars_repo_name: moritz-g/ExercisesRepo

max_issues_repo_name: moritz-g/ExercisesRepo

max_forks_repo_name: moritz-g/ExercisesRepo


text:
type MLEst
    rho
    transform

    MLEst() = new()
end

function MLEst(snr; transform = (x,_) -> x)
    est = MLEst()
    est.rho = 10^(0.1*snr);
    est.transform = transform
    est
end

function estimate(est::MLEst, y)
    z = est.transform(y, :notransp)
    (nAntennas,nCoherence,nBatches) = size(z)# input data in transformed domain

    cest = max.( 0, sum(abs2,z,2) - nCoherence/est.rho )
    cest = reshape(cest,nAntennas,nBatches)

    W = cest./(cest .+ nCoherence/est.rho)
    W = reshape(W,nAntennas,1,nBatches)

    hest = est.transform( W .* z, :transp)
end

max_stars_repo_name: tum-msv/learning-mmse-est


text:
################################################################################
#
#	NEW DEFINITIONS OF LATTICES AND UNITCELLS
#
################################################################################

# usings
using LinearAlgebra

################################################################################
#
#	ABSTRACT TYPES
#
################################################################################

# Bonds
abstract type AbstractBond{L,N} end

# Sites
abstract type AbstractSite{L,D} end

# Unitcells
abstract type AbstractUnitcell{D,N,L,B<:AbstractBond{L,N},S<:AbstractSite{L,D}} end

# Lattices
abstract type AbstractLattice{D,N,L,B<:AbstractBond{L,N},S<:AbstractSite{L,D},U<:AbstractUnitcell{D,N,L,B,S}} end

################################################################################
#
#	INTERFACING / ACCESSING
#	ABSTRACT TYPES
#	(have to be overwritten by concrete types)
#
################################################################################

# BONDS

# default constructor interface
function newBond(index_from::Int64, index_to::Int64, label::L, wrap::NTuple{N,Int64}, ::Type{B}) :: B where {L,N,B<:AbstractBond{L,N}} error("not implemented function 'newBond' for bond type $(B) with label type $(L) and wrap length $(N)") end

# from and to index
function indexFrom(b::AbstractBond{L,N}) :: Int64 where {L,N} error("not implemented interface function 'indexFrom' for bond type $(typeof(b))") end
function indexTo(b::AbstractBond{L,N})   :: Int64 where {L,N} error("not implemented interface function 'indexTo' for bond type $(typeof(b))") end

# label
function label(b::AbstractBond{L,N}) :: L where {L,N} error("not implemented interface function 'label' for bond type $(typeof(b))") end

# wrap
function wrap(b::AbstractBond{L,N}) :: NTuple{N,Int64} where {L,N} error("not implemented interface function 'wrap' for bond type $(typeof(b))") end

# TESTING THE BOND INTERFACE
function testInterface(::Type{T}) where {T<:AbstractBond}
	# get the parameterless constructor
	B = Base.typename(T).wrapper
	# get a new bond
	bond_1 = newBond(1, 1, 1.0, (1,1), B{Float64, 2})
	bond_2 = newBond(1, 1, "tx", (1,1), B{String, 2})
	bond_3 = newBond(1, 1, 2, (1,1,1), B{Int64, 3})
	# get the indices and labels and wraps
	for b in [bond_1, bond_2, bond_3]
		indexFrom(b)
		indexTo(b)
		label(b)
		wrap(b)
	end
	# return nothing
	return nothing
end

# SITES

# default constructor interface
function newSite(position::Vector{Float64}, label::L, ::Type{S}) :: S where {L,S<:AbstractSite{L,D} where D} error("not implemented function 'newSite' for site type $(S) with label type $(L)") end

# position
function point(s::AbstractSite{L,D}) :: Vector{Float64} where {L,D} error("not implemented interface function 'point' for site type $(typeof(s))") end

# label
function label(s::AbstractSite{L,D}) :: L where {L,D} error("not implemented interface function 'label' for site type $(typeof(s))") end

# TESTING THE SITE INTERFACE
function testInterface(::Type{T}) where {T<:AbstractSite}
	# get the parameterless constructor
	S = Base.typename(T).wrapper
	# get a new site
	site_1 = newSite([1.0, 0.0], "site_1", S{String,2})
	site_2 = newSite([1.0, 0.0, 0.0], 1.0, S{Float64,3})
	site_3 = newSite([1.0, 0.0], 1, S{Int64,2})
	# get the points and labels
	for s in [site_1, site_2, site_3]
		point(s)
		label(s)
	end
	# return nothing
	return nothing
end

# DEFINITION OF A BOND
# L - Label type
# N - dimension of the Bravais lattice (i.e. number of primitive directions to wrap)
mutable struct Bond{L,N} <: AbstractBond{L,N}

	# indices that the bond connects
	from	::Int64
	to		::Int64

	# label
	label	::L

	# wrap
	wrap	::NTuple{N, Int64}

end

# IMPLEMENTATION OF BOND RELATED FUNCTIONS

# creation
function newBond(index_from::Int64, index_to::Int64, label::L, wrap::NTuple{N,Int64}, ::Type{Bond{L,N}}) :: Bond{L,N} where {L,N} return Bond{L,N}(index_from, index_to, label, wrap) end

# from and to index
function indexFrom(b::Bond{L,N}) :: Int64 where {N,L} return b.from end
function indexTo(b::Bond{L,N})   :: Int64 where {N,L} return b.to end

# label
function label(b::Bond{L,N}) :: L where {N,L} return b.label end

# wrap
function wrap(b::Bond{L,N}) :: NTuple{N,Int64} where {N,L} return b.wrap end

# DEFINITION OF A SITE
# L - Label type
# D - dimension of the real space in which the site is located
mutable struct Site{L,D} <: AbstractSite{L,D}

	# label
	label	:: L

	# point
	point	:: Vector{Float64}

end

# IMPLEMENTATION OF SITE RELATED FUNCTIONS

# default constructor interface
function newSite(position::Vector{Float64}, label::L, ::Type{Site{L,D}}) :: Site{L,D} where {L,D} return Site{L,D}(label, position) end

# position
function point(s::Site{L,D}) :: Vector{Float64} where {L,D} return s.point end

# label
function label(s::Site{L,D}) :: L where {L,D} return s.label end

# DEFINITION OF A UNITCELL
# D - space dimension
# N - dimension of the Bravais lattice
# L - label type
mutable struct Unitcell{D,N,L,B,S} <: AbstractUnitcell{D,N,L,B,S}

    # FIELDS

    # basis vectors of the Bravais lattice
    lattice_vectors	::Vector{Vector{Float64}}

    # basis sites within the unitcell
    sites			::Vector{S}

    # list of bonds
    bonds			::Vector{B}

end

# Unitcell obtaining (2D - Square)
function getUnitcellSquare_New()::Unitcell_New{2,2,Float64}
	# return the unitcell
	return Unitcell_New{2,2,Float64}(
		# Bravais lattice vectors
	 	Vector{Float64}[
	    	[1.0, 0.0],
	    	[0.0, 1.0]
		],
	    # Basis Definition
	    Vector{Float64}[
	        [0.0, 0.0]
	    ],
	    # Bond Definition
	    Bond{Float64,2}[
	        Bond(1,1, 1.0, ( 0, 1) ),
	        Bond(1,1, 1.0, ( 0,-1) ),
	        Bond(1,1, 1.0, ( 1, 0) ),
	        Bond(1,1, 1.0, (-1, 0) )
	    ]
	)
end

# Unitcell obtaining (3D - SC)
function getUnitcellSC_New()::Unitcell_New{3,3,Float64}
	# return the unitcell
	return Unitcell_New{3,3,Float64}(
		# Bravais lattice vectors
		Vector{Float64}[
	        [1.0, 0.0, 0.0],
	        [0.0, 1.0, 0.0],
	        [0.0, 0.0, 1.0]
		],
        # Basis Definition
        Vector{Float64}[
            [0.0, 0.0, 0.0]
        ],
		# Bond Definition
	    Bond{Float64,3}[
            Bond(1, 1, 1.0, ( 1, 0, 0)),
            Bond(1, 1, 1.0, (-1, 0, 0)),
            Bond(1, 1, 1.0, (0,  1, 0)),
            Bond(1, 1, 1.0, (0, -1, 0)),
            Bond(1, 1, 1.0, (0, 0,  1)),
            Bond(1, 1, 1.0, (0, 0, -1)),
        ]
	)
end

# DEFINITION OF A LATTICE
# D - space dimension
# N - dimension of the Bravais lattice
# L - label type
mutable struct Lattice_New{D,N,L}

    # FIELDS

	# The unitcell of the lattice if applicable
    unitcell			::Unitcell_New{D,N,L}
	# how often the unitcell is repeated in the elementary directions
    unitcellRepetitions	::Array{Int64, 1}

    # basis vectors of the Bravais lattice
    lattice_vectors		::Array{Vector{Float64}, 1}

    # basis sites within the unitcell
    sites				::Array{Vector{Float64}, 1}
	# array of ints that give the index of the site in the original unit cell
    site_indices		::Array{Int64, 1}

    # list of bonds
    bonds				::Array{Bond{L,N}, 1}

end

# define the index function to get the index of an element in the array
function index(
			i1::Int64, i2::Int64,
			N1::Int64, N2::Int64
		) :: Int64
	return (i1-1)*N2 + i2
end
function index(
			i1::Int64, i2::Int64, i3::Int64,
			N1::Int64, N2::Int64, N3::Int64
		) :: Int64
	return N3*((i1-1)*N2+i2-1) + i3
end
function index(
			i1::Int64, i2::Int64, i3::Int64, i4::Int64,
			N1::Int64, N2::Int64, N3::Int64, N4::Int64
		) :: Int64
	return N4*(N3*((i1-1)*N2+i2-1) + i3 -1) + i4
end

# function to constrain an index
# returns (index_constrained, offset)
# s.t. index = index_constraint + offset*index_max
# To be replaced by ic = div(.,.) and offset = rem(.,.) ???
# apparently not faster...
function getConstrainedIndex(
		index		::Int64,
		index_max	::Int64
	) :: Tuple{Int64, Int64}
	# specify the initial offset as 0
	offset = 0 :: Int64
	# check that the index is bigger than the minimum value
	while index < 1
		index += index_max
		offset += -1
	end
	# check that the index is smaller than the maximum value
	while index > index_max
		index -= index_max
		offset += 1
	end
	# return both index and offset
	return (index, offset)
end

# Construction of periodic lattice (in 2D)
function getLatticePeriodic_New(
			unitcell			::Unitcell_New{D,2,L},
			repetition_array	::Array{Int64, 1}
		) :: Lattice_New{D,2,L} where {D,L}

    # extract the cardinal directions of the lattice from the array
    N_a1	= abs(repetition_array[1]) 	:: Int64
    N_a2	= abs(repetition_array[2]) 	:: Int64
	N_sites = length(unitcell.sites)	:: Int64
	N_bonds = length(unitcell.bonds)	:: Int64

    # GENERATE NEW POSITIONS
	sites 				= Array{Vector{Float64},1}(undef, N_a1*N_a2*N_sites)
    site_indices 		= Array{Int64,1}(undef, N_a1*N_a2*N_sites)

	# set all positions to their correct values
	for i in 1:N_a1
	for j in 1:N_a2
	    @simd for a in 1:N_sites
			# calculate the site index
			site_index 					= index(i,j,a, N_a1,N_a2,N_sites)
			# calculate the position
			sites[site_index]			= unitcell.sites[a] .+ (i.*unitcell.lattice_vectors[1]) .+ (j.*unitcell.lattice_vectors[2])
			# set the site index
	        site_indices[site_index]	= a
		end
	end
    end

    # GENERATE NEW CONNECTIONS
	bonds = Array{Bond{L,2},1}(undef, N_a1*N_a2*N_bonds)

	# iterate over all unit cells
	for i in 1:N_a1
	for j in 1:N_a2
		# add all connections for unitcell (i,j)
		@simd for b in 1:N_bonds
			# get the respective bond
			bond = unitcell.bonds[b] :: Bond{L,2}
			# calculate the index from where the connection goes
			index_from = index(i,j,bond.from, N_a1,N_a2,N_sites)
			# calculate the aimed unitcell
			i_to = i + bond.wrap[1]
			j_to = j + bond.wrap[2]
			# check if the connection goes around in a1 direction
            (i_to, offset_a1) = getConstrainedIndex(i_to, N_a1)
			# check if the connection goes around in a2 direction
            (j_to, offset_a2) = getConstrainedIndex(j_to, N_a2)
			# get the index to where the connection goes
			index_to = index(i_to, j_to, bond.to, N_a1,N_a2,N_sites)
            # generate a new connection and set it in the list
			bonds[index(i,j,b, N_a1,N_a2,N_bonds)] = Bond{L,2}(index_from, index_to, bond.label, (offset_a1, offset_a2))
		end
	end
	end

    # generate new lattice vectors, now spanning the WHOLE lattice
    lattice_vectors = Vector{Float64}[
			# multiply by extent of lattice
			unitcell.lattice_vectors[1] .* N_a1,
			unitcell.lattice_vectors[2] .* N_a2
		]

    # save everything to a Lattice object
    lattice = Lattice_New{D,2,L}(
        	unitcell,
        	Int64[N_a1, N_a2],
        	lattice_vectors,
        	sites,
        	site_indices,
        	bonds
        )

    # return the lattice
    return lattice

end

# Construction of periodic lattice (in 3D)
function getLatticePeriodic_New(
			unitcell			::Unitcell_New{D,3,L},
			repetition_array	::Array{Int64, 1}
		) :: Lattice_New{D,3,L} where {D,L}

    # extract the cardinal directions of the lattice from the array
    N_a1	= abs(repetition_array[1]) 	:: Int64
    N_a2	= abs(repetition_array[2]) 	:: Int64
    N_a3	= abs(repetition_array[3]) 	:: Int64
	N_sites = length(unitcell.sites)	:: Int64
	N_bonds = length(unitcell.bonds)	:: Int64

    # GENERATE NEW POSITIONS
	sites 				= Array{Vector{Float64},1}(undef, N_a1*N_a2*N_a3*N_sites)
    site_indices 		= Array{Int64,1}(undef, N_a1*N_a2*N_a3*N_sites)

	# set all positions to their correct values
	for i in 1:N_a1
	for j in 1:N_a2
	for k in 1:N_a3
	    @simd for a in 1:N_sites
			# calculate the site index
			site_index 					= index(i,j,k,a, N_a1,N_a2,N_a3,N_sites)
			# calculate the position
			sites[site_index]			= unitcell.sites[a] .+ (i.*unitcell.lattice_vectors[1]) .+ (j.*unitcell.lattice_vectors[2]) .+ (k.*unitcell.lattice_vectors[3])
			# set the site index
	        site_indices[site_index]	= a
		end
	end
	end
    end

    # GENERATE NEW CONNECTIONS
	bonds = Array{Bond{L,3},1}(undef, N_a1*N_a2*N_a3*N_bonds)

	# iterate over all unit cells
	for i in 1:N_a1
	for j in 1:N_a2
	for k in 1:N_a3
		# add all connections for unitcell (i,j)
		@simd for b in 1:N_bonds
			# get the respective bond
			bond = unitcell.bonds[b] :: Bond{L,3}
			# calculate the index from where the connection goes
			index_from = index(i,j,k,bond.from, N_a1,N_a2,N_a3,N_sites)
			# calculate the aimed unitcell
			i_to = i + bond.wrap[1] ::Int64
			j_to = j + bond.wrap[2] ::Int64
			k_to = k + bond.wrap[3] ::Int64
			# check if the connection goes around in a1 direction
            (i_to, offset_a1) = getConstrainedIndex(i_to, N_a1)
			# check if the connection goes around in a2 direction
            (j_to, offset_a2) = getConstrainedIndex(j_to, N_a2)
			# check if the connection goes around in a3 direction
            (k_to, offset_a3) = getConstrainedIndex(k_to, N_a3)
			# get the index to where the connection goes
			index_to = index(i_to, j_to, k_to, bond.to, N_a1,N_a2,N_a3,N_sites)
            # generate a new connection and set it in the list
			bonds[index(i,j,k,b, N_a1,N_a2,N_a3,N_bonds)] = Bond{L,3}(index_from, index_to, bond.label, (offset_a1, offset_a2, offset_a3))
		end
	end
	end
	end

	# generate new lattice vectors, now spanning the WHOLE lattice
    lattice_vectors = Vector{Float64}[
			# multiply by extent of lattice
			unitcell.lattice_vectors[1] .* N_a1,
			unitcell.lattice_vectors[2] .* N_a2,
			unitcell.lattice_vectors[3] .* N_a3
		]

    # save everything to a Lattice object
    lattice = Lattice_New{D,3,L}(
        	unitcell,
        	Int64[N_a1, N_a2, N_a3],
        	lattice_vectors,
        	sites,
        	site_indices,
        	bonds
        )

    # return the lattice
    return lattice

end


text:
"""
    VerticallyStretchedCartesianGrid{FT, TX, TY, TZ, R, A} <: AbstractGrid{FT, TX, TY, TZ}

A Cartesian grid with with constant horizontal grid spacings `Δx` and `Δy`, and
non-uniform or stretched vertical grid spacing `Δz` between cell centers and cell faces,
topology `{TX, TY, TZ}`, and coordinate ranges of type `R` (where a range can be used) and
`A` (where an array is needed).
"""
struct VerticallyStretchedCartesianGrid{FT, TX, TY, TZ, R, A} <: AbstractGrid{FT, TX, TY, TZ}
    # Number of grid points in (x,y,z).
    Nx::Int
    Ny::Int
    Nz::Int
    # Halo size in (x,y,z).
    Hx::Int
    Hy::Int
    Hz::Int
    # Domain size [m].
    Lx::FT
    Ly::FT
    Lz::FT
    # Grid spacing [m].
    Δx::FT
    Δy::FT
    ΔzF::A
    ΔzC::A
    # Range of coordinates at the centers of the cells.
    xC::R
    yC::R
    zC::A
    # Range of grid coordinates at the faces of the cells.
    # Note: there are Nx+1 faces in the x-dimension, Ny+1 in the y, and Nz+1 in the z.
    xF::R
    yF::R
    zF::A
end

function VerticallyStretchedCartesianGrid(FT=Float64, arch=CPU();
        size, halo=(1, 1, 1), topology=(Periodic, Periodic, Bounded),
        length=nothing, x=nothing, y=nothing, z=nothing, zF=nothing)

    # Hack that allows us to use `size` and `length` as keyword arguments but then also
    # use the `size` and `length` functions.
    sz, len = size, length
    length = Base.length

    TX, TY, TZ = validate_topology(topology)
    Lx, Ly, Lz, x, y, z = validate_grid_size_and_length(sz, len, halo, x, y, z)

    Nx, Ny, Nz = sz
    Hx, Hy, Hz = halo

    Δx = convert(FT, Lx / Nx)
    Δy = convert(FT, Ly / Ny)

    x₁, x₂ = convert.(FT, [x[1], x[2]])
    y₁, y₂ = convert.(FT, [y[1], y[2]])
    z₁, z₂ = convert.(FT, [z[1], z[2]])

    xF = range(x₁, x₂; length=Nx+1)
    yF = range(y₁, y₂; length=Ny+1)

    xC = range(x₁ + Δx/2, x₂ - Δx/2; length=Nx)
    yC = range(y₁ + Δy/2, y₂ - Δy/2; length=Ny)

    zF, zC, ΔzF, ΔzC = validate_and_generate_variable_grid_spacing(FT, zF, Nz, z₁, z₂)

    VerticallyStretchedCartesianGrid{FT, typeof(TX), typeof(TY), typeof(TZ), typeof(xF), typeof(zF)}(
        Nx, Ny, Nz, Hx, Hy, Hz, Lx, Ly, Lz, Δx, Δy, ΔzF, ΔzC, xC, yC, zC, xF, yF, zF)
end


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

mutable struct Node
    original::Union{Vector{UInt8}, Nothing}
    feature::Union{Vector{UInt8}, Nothing}
    pos::Int32
    epos::Int32
    index::Int32
    left_id::Int32
    right_id::Int32
    cost::Int32
    min_cost::Int32
    back_pos::Int32
    back_index::Int32
    skip::Bool
end

function copy(node::Node)::Node
    original::Union{Vector{UInt8}, Nothing} = nothing
    if node.original != nothing
        original = zeros(UInt8, length(node.original))
        copyto!(original, node.original)
    end

    feature::Union{Vector{UInt8}, Nothing} = nothing
    if node.feature != nothing
        feature = zeros(UInt8, length(node.feature))
        copyto!(feature, node.feature)
    end

    Node(
        original,
        feature,
        node.pos,
        node.epos,
        node.index,
        node.left_id,
        node.right_id,
        node.cost,
        node.min_cost,
        node.back_pos,
        node.back_index,
        node.skip
    )
end

function new_bos()::Node
    Node(
        nothing,    # original
        nothing,    # feature
        0,          # pos
        1,          # epos
        0,          # index
        -1,         # left_id
        0,          # right_id
        0,          # cost
        0,          # min_cost
        -1,         # back_pos
        -1,         # back_index
        false       # skip
    )
end

function new_eos(pos::Int32)::Node
    Node(
        nothing,    # original
        nothing,    # feature
        pos,        # pos
        pos+1,      # epos
        0,          # index
        0,          # left_id
        -1,         # right_id,
        0,          # cost
        0x7FFFFFFF, # min_cost
        -1,         # back_pos
        -1,         # back_index
        false       # skip
    )
end

function new_node(e::DicEntry)
    Node(
        e.original, # entry
        e.feature,  # feature
        0,          # pos
        0,          # epos
        e.posid,    # index
        e.lc_attr,  # left_id
        e.rc_attr,  # right_id,
        e.wcost,    # cost
        0x7FFFFFFF, # min_cost
        -1,         # back_pos
        -1,         # back_index
        e.skip      # skip
    )
end

function is_bos(node::Node)::Bool
    node.original == nothing && node.pos == 0
end

function is_eos(node::Node)::Bool
    node.original == nothing && node.pos != 0
end

function node_len(node::Node)::Int32
    if node.original == nothing
        Int32(1)
    else
        Int32(length(node.original))
    end
end

function dump_node(node::Node)
    if node.original != nothing
        original = String(node.original)
    else
        original = ""
    end
    if node.feature != nothing
        feature = String(node.feature)
    else
        feature = ""
    end
    println("$(original),$(feature),$(node_len(node)),$(node.pos),$(node.epos),$(node.index),$(node.left_id),$(node.right_id),$(node.cost),$(node.min_cost),$(node.back_pos),$(node.back_index),$(node.skip)")
end

function dump_nodes_list(prompt::AbstractString, nodes_list::Vector{Vector{Node}})
    println("++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++")
    for nodes in nodes_list
        println("--------------------------------------------------------------------------------------------------------------")
        dump_nodes(nodes)
    end
end

function dump_nodes(nodes::Vector{Node})
    println("--------------------------------------------------------------------------------------------------------------")
    for node in nodes
        dump_node(copy(node))
    end
end

mutable struct Lattice
    snodes::Vector{Vector{Node}}
    enodes::Vector{Vector{Node}}
    p::Int32
end

function new_lattice(size::Int64)
    snodes::Vector{Vector{Node}} = []
    enodes::Vector{Vector{Node}} = []

    push!(enodes,[])
    for _ in 1:(size + 2)
        push!(snodes, [])
        push!(enodes, [])
    end
    bos::Node = new_bos()
    push!(snodes[1], bos)
    push!(enodes[2], bos)

    Lattice(snodes, enodes, 1)
end

function add!(lattice::Lattice, node::Node, matrix::Matrix)
    min_cost = node.min_cost
    best_node = lattice.enodes[lattice.p+1][1]

    for enode in lattice.enodes[lattice.p+1]
        if enode.skip
            for enode2 in lattice.enodes[enode.pos+1]
                cost = enode2.min_cost + get_trans_cost(matrix, UInt16(enode2.right_id), UInt16(node.left_id))
                if cost < min_cost
                    min_cost = cost
                    best_node = enode
                end
            end
        else
            cost = enode.min_cost + get_trans_cost(matrix, UInt16(enode.right_id), UInt16(node.left_id))
            if cost < min_cost
                min_cost = cost
                best_node = enode
            end
        end
    end

    node.min_cost = min_cost + node.cost
    node.back_index = best_node.index
    node.back_pos = best_node.pos
    node.pos = lattice.p
    node.epos = lattice.p + node_len(node)
    node.index = length(lattice.snodes[lattice.p+1])
    push!(lattice.snodes[node.pos+1], deepcopy(node))
    push!(lattice.enodes[node.epos+1], deepcopy(node))
end

function forward(lattice::Lattice)::Int64
    old_p = lattice.p
    lattice.p += 1
    while length(lattice.enodes[lattice.p+1]) == 0
        lattice.p += 1
    end
    lattice.p - old_p
end

function end!(lattice::Lattice, matrix::Matrix)
    add!(lattice, new_eos(lattice.p), matrix)
    lattice.snodes = lattice.snodes[1:lattice.p+1]
    lattice.enodes = lattice.enodes[1:lattice.p+2]
end

function backward(lattice::Lattice)::Vector{Node}
    @assert is_eos(lattice.snodes[length(lattice.snodes)][1])

    shortest_path::Vector{Node} = []
    pos = length(lattice.snodes) -1
    index = 0
    while pos >= 0
       node = lattice.snodes[pos+1][index+1]
       index = node.back_index
       pos = node.back_pos
       push!(shortest_path, node)
    end
    reverse(shortest_path)
end

function backward_astar(lattice::Lattice, n::Number, matrix::Matrix)::Vector{Vector{Node}}
    paths::Vector{Vector{Node}} = []
    epos = length(lattice.enodes) -1
    node = lattice.enodes[epos+1][1]
    @assert is_eos(node)

    pq = PriorityQueue{BackwardPath, Int32}()
    bp = BackwardPath(matrix, deepcopy(node), nothing)
    enqueue!(pq, bp, bp.cost_from_bos + bp.cost_from_eos)

    while length(pq) > 0 && n > 0
        bp = dequeue!(pq)
        if is_complete(bp)
            push!(paths, reverse(bp.back_path))
            n -= 1
        else
            node = bp.back_path[length(bp.back_path)]
            epos = node.epos - node_len(node)
            for i in 1:(length(lattice.enodes[epos+1]))
                node = lattice.enodes[epos+1][i]
                bp2 = BackwardPath(matrix, deepcopy(node), bp)
                enqueue!(pq, bp2, bp2.cost_from_bos + bp2.cost_from_eos)
            end
        end
    end

    paths
end

struct BackwardPath
    cost_from_bos::Int32
    cost_from_eos::Int32
    back_path::Vector{Node}

    function BackwardPath(matrix::Matrix, node::Node, right_path::Union{BackwardPath, Nothing})
        cost_from_bos = node.min_cost
        cost_from_eos = 0
        back_path::Vector{Node} = []

        if right_path != nothing
            neighbor_node = right_path.back_path[length(right_path.back_path)]
            cost_from_eos = right_path.cost_from_eos + neighbor_node.cost + get_trans_cost(matrix, UInt16(node.right_id), UInt16(neighbor_node.left_id))
            for node in right_path.back_path
                push!(back_path, deepcopy(node))
            end
        else
            @assert is_eos(node)
        end

        push!(back_path, deepcopy(node))

        new(cost_from_bos, cost_from_eos, back_path)
    end
end

function is_complete(bp::BackwardPath)::Bool
    is_bos(bp.back_path[length(bp.back_path)])
end


text:
"""
    initialenv(mps, mpo)

construct and initialize the environment for a dmrg algorithm that is
about to start at positon `at` (defaulted to `1`)! The environment is
made from `mps` and `mpo` by the _mpsupdateright function.

"""
function initialenv(mps::MPState{Ys},
                    mpo::MPOperator{Yo};
                    at::Int=1) where {Ys, Yo}
    vtype(Ys) == vtype(Yo) && eltype(Ys) == eltype(Yo) ||
        error("MPS, MPO not the same type!")
    center(mps) == at || error("MPS center has to be $at")

    lx = length(mps)
    env = Vector{Ys}(undef, lx+2)

    lVa = leftspace(mps)
    rVa = rightspace(mps)
    lVw = leftspace(mpo)
    rVw = rightspace(mpo)

    T = eltype(Ys)
    env[1] = fill(one(T), (dual(lVa), dual(lVw), lVa))
    env[lx+2] = fill(one(T), (rVa, rVw, dual(rVa)))

    for l = lx:-1:at
        env[l+1] = _mpsupdateright(env[l+2], mps.As[l], mpo.Ws[l])
    end

    for l = 1:at-1
        env[l+1] = _mpsupdateleft(env[l], mps.As[l], mpo.Ws[l])
    end
    return env
end

# function _initialenv(mps::MPState{Y1},
#                      mpo::MPOperator{Y2};
#                      at::Int=1) where {Y1, Y2}
#     Y = promote_type(Y1, Y2)
#     initialenv(convert(MPState{Y}, mps),
#                convert(MPOperator{Y}, mpo),
#                at=at)
# end

"""
    dmrg2sitesweep!(mps, mpo, env; [maxdim, tol, verbose, kwargs...])

Performs a single sweep (left to right and back) of 2site DMRG on
`mps` where the hamiltonian is given in `mpo` and environment matrices
for each step are stored in `env`. The singular values are truncated
upto truncation error or up to number of `maxdim`.

Returns the tuple of (energies, truncations) and updates, `mps` and `env`
"""

function dmrg2sitesweep!(mps::MPState{Ys},
                         mpo::MPOperator{Yo},
                         env::Vector{Ys};
                         maxdim::Int=200,
                         tol::Float64=1.e-9,
                         lanczostol::Float64=1.e-7,
                         krylovdim::Int=5,
                         krylovmaxiter::Int=8,
                         verbose::Bool=false) where {Ys, Yo}

    vtype(Ys) == vtype(Yo) && eltype(Ys) == eltype(Yo) ||
        error("MPS, MPO not the same type!")
    lx = length(mps)
    mps.center == 1 || error("The center of MPS has to be 1 for dmrg L->R->L")

    energies = Vector{Float64}()
    A = mps.As[1]
    for l = 1:lx-2
        AA = contract(A, (1,2,-1), mps.As[l+1], (-1,3,4))
        es, vs, info = eigsolve(v->_applymps2site(v, env[l], env[l+3],
                                                  mpo.Ws[l], mpo.Ws[l+1]),
                                AA, 1, :SR, ishermitian=true,
                                tol=lanczostol,
                                krylovdim=krylovdim,
                                maxiter=krylovmaxiter)
        v = vs[1]
        #e = es[1]
        verbose && println("Sweep L2R: bond $l, $(l+1) -> energy $(es[1])")
        verbose && println("normres = $(info.normres[1]), #iterations = $(info.numiter), #applications = $(info.numops)")

        u,s,v = svdtrunc(SymMatrix(v, [1,2], [3,4]), maxdim=maxdim, tol=tol)
        verbose && println("truncation error = $(1-norm(s))")
        normalize!(s)
        AA = splitleg(splitleg(u*s*v, 2, space(AA)[3:4]), 1, space(AA)[1:2])
        e = dot(AA, _applymps2site(AA, env[l], env[l+3],
                                   mpo.Ws[l], mpo.Ws[l+1]))
        verbose && println("energy after truncation $e")
        push!(energies, e)
        if verbose
            values = diag(s)
            println("Entanglement S1 = $(entropy(values.^2))")
            println()
        end
        mps.As[l] = splitleg(u, 1, space(A)[1:2])

        env[l+1] = _mpsupdateleft(env[l], mps.As[l], mpo.Ws[l])

        A = splitleg(s*v, 2, space(AA)[3:4])
    end

    l = lx-1
    AA = contract(A, (1,2,-1), mps.As[l+1], (-1,3,4))
    es, vs, info = eigsolve(v->_applymps2site(v, env[l], env[l+3],
                                              mpo.Ws[l], mpo.Ws[l+1]),
                            AA, 1, :SR, ishermitian=true,
                            tol=lanczostol,
                            krylovdim=krylovdim,
                            maxiter=krylovmaxiter)

    v = vs[1]
    #e = es[1]
    verbose && println("Sweep L2R: bond $l, $(l+1) -> energy $(es[1])")
    verbose && println("normres = $(info.normres[1]), #iterations = $(info.numiter), #applications = $(info.numops)")
    # verbose && println("truncation error = 1")
    # push!(energies, e)

    for l = lx-1:-1:2
        u,s,v = svdtrunc(SymMatrix(v, [1,2], [3,4]), maxdim=maxdim, tol=tol)
        verbose && println("truncation error = $(1-norm(s))")
        normalize!(s)
        AA = splitleg(splitleg(u*s*v, 2, space(AA)[3:4]), 1, space(AA)[1:2])
        e = dot(AA, _applymps2site(AA, env[l], env[l+3],
                                   mpo.Ws[l], mpo.Ws[l+1]))
        verbose && println("energy after truncation $e")
        push!(energies, e)

        if verbose
            values = diag(s)
            println("Entanglement S1 = $(entropy(values.^2))")
            println()
        end

        mps.As[l+1] = splitleg(v, 2, space(AA)[3:4])
        env[l+2] = _mpsupdateright(env[l+3], mps.As[l+1], mpo.Ws[l+1])

        A = splitleg(u*s, 1, space(AA)[1:2])
        AA = contract(mps.As[l-1], (1,2,-1), A, (-1,3,4))

        es, vs, info = eigsolve(v->_applymps2site(v, env[l-1], env[l+2],
                                                  mpo.Ws[l-1], mpo.Ws[l]),
                                AA, 1, :SR, ishermitian=true,
                                tol=lanczostol,
                                krylovdim=krylovdim,
                                maxiter=krylovmaxiter)
        v = vs[1]
        #e = es[1]
        verbose && println("Sweep R2L: bond $(l-1), $l -> energy $(es[1])")
        verbose && println("normres = $(info.normres[1]), #iterations = $(info.numiter), #applications = $(info.numops)")
        push!(energies, e)
    end
    l = 1
    u,s,v = svdtrunc(SymMatrix(v, [1,2], [3,4]), maxdim=maxdim, tol=tol)
    verbose && println("truncation error = $(1-norm(s))")
    normalize!(s)
    AA = splitleg(splitleg(u*s*v, 2, space(AA)[3:4]), 1, space(AA)[1:2])
    e = dot(AA, _applymps2site(AA, env[l], env[l+3],
                               mpo.Ws[l], mpo.Ws[l+1]))
    verbose && println("energy after truncation $e")
    push!(energies, e)

    if verbose
        values = diag(s)
        println("Entanglement S1 = $(entropy(values.^2))")
        println()
    end

    mps.As[l+1] = splitleg(v, 2, space(AA)[3:4])
env[l+2] = _mpsupdateright(env[l+3], mps.As[l+1], mpo.Ws[l+1])

mps.As[l] = splitleg(u*s, 1, space(AA)[1:2])
return energies
end

# function dmrg2site!(mps::MPState{S,Tv},
#                     mpo::MPOperator{Tv};
#                     infostring::String="",
#                     finalmeasurements::Vector{Measurement}=[],
#                     sweepmeasurements::Vector{Measurement}=[],
#                     n_sweeps::Int=4,
#                     maxdim::Int=200,
#                     tol::Float64=1.e-8,
#                     verbose::Bool=false)

#     n < 1 && return

#     env = initialenv(mps, mpo)
#     energies = zeros(Float64, 2*mps.lx-3, n_sweeps)
#     for n in 1:n_sweeps
#         verbose && println("Starting sweeps ")
#         energies[:, n] =
#             dmrg2sitesweep!(mps, mpo, env, maxdim=maxdim, tol=tol, verbose=verbose)
#         if length(sweepmeasurements) > 0
#             # do the measurements
#         end
#     end

# end

max_stars_repo_name: amiragha/mps

max_issues_repo_name: amiragha/mps

max_forks_repo_name: amiragha/mps


text:
module CSGBuilding

using AbstractTrees
using LinearAlgebra: dot, norm, normalize, eigvecs, cross
using StaticArrays: SVector, SMatrix
using Random
using Parameters
using Logging
#import Base.Threads.@spawn
using Base.Threads
using GeometryTypes: normals, vertices
using MeshIO
using FileIO: load
using Base: Semaphore, acquire, release
# cone
using RANSAC: project2cone
# translational
using RANSAC: project2sketchplane, dn2shape_outw, outwardsnormal
# conversion
using RANSAC: FittedPlane, FittedSphere, FittedCylinder, FittedCone, ExtractedTranslational
using RANSAC: PointCloud, ScoredShape, ShapeCandidate
# others
using RANSAC: findAABB
# visualize and write
import JSON

export  ImplicitResult,
        value,
        normal,
        ImplicitSphere,
        ImplicitPlane,
        ImplicitCylinder,
        ImplicitCone,
        ImplicitTranslational,
        evaluate

export  CSGNode,
        valueandnormal,
        depth,
        treesize,
        evaldistance,
        randomtree

export  CachedSurface,
        CachedResult,
        CachedCSGNode,
        cachenodes,
        randomcachedtree,
        cached2normaltree

export  tree2func

export  writeparaviewformat,
        toNiceTree,
        toJSON,
        rmse,
        readobj

export  centroid,
        findOBB,
        issame,
        ransacresult2implicit

export  CSGGeneticBuildParameters,
        rankpopulation,
        crossover,
        mutate,
        geneticbuildtree,
        cachedgeneticbuildtree,
        cachedfuncgeneticbuildtree

export  toimplicit,
        scored2implicit

export  tree2tex,
        tree2tex_orderfix,
        json2tex

export  isosurface,
        writeOBJ

include("implicitsurfaces.jl")
include("csgtrees.jl")
include("cachedtrees.jl")
include("operations.jl")
include("codegen.jl")
include("visualize.jl")
include("orientedbox.jl")
include("genetic.jl")
include("deprecated.jl")
include("ransaccompat.jl")
include("textree.jl")
include("dual_contour.jl")

end # module


text:
int_rules_4_5_11 = @theory begin

    #= ::Subsection::Closed:: =#
    #= 4.5.11*(e*x)^m*(a+b*sec(c+d*x^n))^p =#
    @apply_utils Antiderivative((~(a') + ~(b') * sec(~(c') + ~(d') * (~x) ^ ~n)) ^ ~(p'), ~x) => (1 / ~n) * Subst(Antiderivative((~x) ^ (1 / ~n - 1) * (~a + ~b * sec(~c + ~d * ~x)) ^ ~p, ~x), ~x, (~x) ^ ~n) <-- FreeQ([~a, ~b, ~c, ~d, ~p], ~x) && (IGtQ(1 / ~n, 0) && IntegerQ(~p))
    @apply_utils Antiderivative((~(a') + ~(b') * csc(~(c') + ~(d') * (~x) ^ ~n)) ^ ~(p'), ~x) => (1 / ~n) * Subst(Antiderivative((~x) ^ (1 / ~n - 1) * (~a + ~b * csc(~c + ~d * ~x)) ^ ~p, ~x), ~x, (~x) ^ ~n) <-- FreeQ([~a, ~b, ~c, ~d, ~p], ~x) && (IGtQ(1 / ~n, 0) && IntegerQ(~p))
    @apply_utils Antiderivative((~(a') + ~(b') * sec(~(c') + ~(d') * (~x) ^ ~n)) ^ ~(p'), ~x) => Unintegrable((~a + ~b * sec(~c + ~d * (~x) ^ ~n)) ^ ~p, ~x) <-- FreeQ([~a, ~b, ~c, ~d, ~n, ~p], ~x)
    @apply_utils Antiderivative((~(a') + ~(b') * csc(~(c') + ~(d') * (~x) ^ ~n)) ^ ~(p'), ~x) => Unintegrable((~a + ~b * csc(~c + ~d * (~x) ^ ~n)) ^ ~p, ~x) <-- FreeQ([~a, ~b, ~c, ~d, ~n, ~p], ~x)
    @apply_utils Antiderivative((~(a') + ~(b') * sec(~(c') + ~(d') * (~u) ^ ~n)) ^ ~(p'), ~x) => (1 / Coefficient(~u, ~x, 1)) * Subst(Antiderivative((~a + ~b * sec(~c + ~d * (~x) ^ ~n)) ^ ~p, ~x), ~x, ~u) <-- FreeQ([~a, ~b, ~c, ~d, ~n, ~p], ~x) && (LinearQ(~u, ~x) && NeQ(~u, ~x))
    @apply_utils Antiderivative((~(a') + ~(b') * csc(~(c') + ~(d') * (~u) ^ ~n)) ^ ~(p'), ~x) => (1 / Coefficient(~u, ~x, 1)) * Subst(Antiderivative((~a + ~b * csc(~c + ~d * (~x) ^ ~n)) ^ ~p, ~x), ~x, ~u) <-- FreeQ([~a, ~b, ~c, ~d, ~n, ~p], ~x) && (LinearQ(~u, ~x) && NeQ(~u, ~x))
    @apply_utils Antiderivative((~(a') + ~(b') * sec(~u)) ^ ~(p'), ~x) => Antiderivative((~a + ~b * sec(ExpandToSum(~u, ~x))) ^ ~p, ~x) <-- FreeQ([~a, ~b, ~p], ~x) && (BinomialQ(~u, ~x) && Not(BinomialMatchQ(~u, ~x)))
    @apply_utils Antiderivative((~(a') + ~(b') * csc(~u)) ^ ~(p'), ~x) => Antiderivative((~a + ~b * csc(ExpandToSum(~u, ~x))) ^ ~p, ~x) <-- FreeQ([~a, ~b, ~p], ~x) && (BinomialQ(~u, ~x) && Not(BinomialMatchQ(~u, ~x)))
    @apply_utils Antiderivative((~x) ^ ~(m') * (~(a') + ~(b') * sec(~(c') + ~(d') * (~x) ^ ~n)) ^ ~(p'), ~x) => (1 / ~n) * Subst(Antiderivative((~x) ^ (Simplify((~m + 1) / ~n) - 1) * (~a + ~b * sec(~c + ~d * ~x)) ^ ~p, ~x), ~x, (~x) ^ ~n) <-- FreeQ([~a, ~b, ~c, ~d, ~m, ~n, ~p], ~x) && (IGtQ(Simplify((~m + 1) / ~n), 0) && IntegerQ(~p))
    @apply_utils Antiderivative((~x) ^ ~(m') * (~(a') + ~(b') * csc(~(c') + ~(d') * (~x) ^ ~n)) ^ ~(p'), ~x) => (1 / ~n) * Subst(Antiderivative((~x) ^ (Simplify((~m + 1) / ~n) - 1) * (~a + ~b * csc(~c + ~d * ~x)) ^ ~p, ~x), ~x, (~x) ^ ~n) <-- FreeQ([~a, ~b, ~c, ~d, ~m, ~n, ~p], ~x) && (IGtQ(Simplify((~m + 1) / ~n), 0) && IntegerQ(~p))
    @apply_utils Antiderivative((~x) ^ ~(m') * (~(a') + ~(b') * sec(~(c') + ~(d') * (~x) ^ ~n)) ^ ~(p'), ~x) => Unintegrable((~x) ^ ~m * (~a + ~b * sec(~c + ~d * (~x) ^ ~n)) ^ ~p, ~x) <-- FreeQ([~a, ~b, ~c, ~d, ~m, ~n, ~p], ~x)
    @apply_utils Antiderivative((~x) ^ ~(m') * (~(a') + ~(b') * csc(~(c') + ~(d') * (~x) ^ ~n)) ^ ~(p'), ~x) => Unintegrable((~x) ^ ~m * (~a + ~b * csc(~c + ~d * (~x) ^ ~n)) ^ ~p, ~x) <-- FreeQ([~a, ~b, ~c, ~d, ~m, ~n, ~p], ~x)
    @apply_utils Antiderivative((~e * ~x) ^ ~(m') * (~(a') + ~(b') * sec(~(c') + ~(d') * (~x) ^ ~n)) ^ ~(p'), ~x) => (((~e) ^ IntPart(~m) * (~e * ~x) ^ FracPart(~m)) / (~x) ^ FracPart(~m)) * Antiderivative((~x) ^ ~m * (~a + ~b * sec(~c + ~d * (~x) ^ ~n)) ^ ~p, ~x) <-- FreeQ([~a, ~b, ~c, ~d, ~e, ~m, ~n, ~p], ~x)
    @apply_utils Antiderivative((~e * ~x) ^ ~(m') * (~(a') + ~(b') * csc(~(c') + ~(d') * (~x) ^ ~n)) ^ ~(p'), ~x) => (((~e) ^ IntPart(~m) * (~e * ~x) ^ FracPart(~m)) / (~x) ^ FracPart(~m)) * Antiderivative((~x) ^ ~m * (~a + ~b * csc(~c + ~d * (~x) ^ ~n)) ^ ~p, ~x) <-- FreeQ([~a, ~b, ~c, ~d, ~e, ~m, ~n, ~p], ~x)
    @apply_utils Antiderivative((~e * ~x) ^ ~(m') * (~(a') + ~(b') * sec(~u)) ^ ~(p'), ~x) => Antiderivative((~e * ~x) ^ ~m * (~a + ~b * sec(ExpandToSum(~u, ~x))) ^ ~p, ~x) <-- FreeQ([~a, ~b, ~e, ~m, ~p], ~x) && (BinomialQ(~u, ~x) && Not(BinomialMatchQ(~u, ~x)))
    @apply_utils Antiderivative((~e * ~x) ^ ~(m') * (~(a') + ~(b') * csc(~u)) ^ ~(p'), ~x) => Antiderivative((~e * ~x) ^ ~m * (~a + ~b * csc(ExpandToSum(~u, ~x))) ^ ~p, ~x) <-- FreeQ([~a, ~b, ~e, ~m, ~p], ~x) && (BinomialQ(~u, ~x) && Not(BinomialMatchQ(~u, ~x)))
    @apply_utils Antiderivative((~x) ^ ~(m') * sec(~(a') + ~(b') * (~x) ^ ~(n')) ^ ~p * sin(~(a') + ~(b') * (~x) ^ ~(n')), ~x) => ((~x) ^ ((~m - ~n) + 1) * sec(~a + ~b * (~x) ^ ~n) ^ (~p - 1)) / (~b * ~n * (~p - 1)) - (((~m - ~n) + 1) / (~b * ~n * (~p - 1))) * Antiderivative((~x) ^ (~m - ~n) * sec(~a + ~b * (~x) ^ ~n) ^ (~p - 1), ~x) <-- FreeQ([~a, ~b, ~p], ~x) && (IntegerQ(~n) && (GeQ(~m - ~n, 0) && NeQ(~p, 1)))
    @apply_utils Antiderivative((~x) ^ ~(m') * csc(~(a') + ~(b') * (~x) ^ ~(n')) ^ ~p * cos(~(a') + ~(b') * (~x) ^ ~(n')), ~x) => (-((~x) ^ ((~m - ~n) + 1)) * csc(~a + ~b * (~x) ^ ~n) ^ (~p - 1)) / (~b * ~n * (~p - 1)) + (((~m - ~n) + 1) / (~b * ~n * (~p - 1))) * Antiderivative((~x) ^ (~m - ~n) * csc(~a + ~b * (~x) ^ ~n) ^ (~p - 1), ~x) <-- FreeQ([~a, ~b, ~p], ~x) && (IntegerQ(~n) && (GeQ(~m - ~n, 0) && NeQ(~p, 1)))
end

max_stars_repo_path: src/rules/4 Trig functions/4.5 Secant/4.5.11 (e x)^m (a+b sec(c+d x^n))^p.jl


text:
function get_file_symbols(thisFile::String)
    return Symbol.(readdlm(show_sourcePath() * "/$thisFile", '\t', String, '\n')[1,:])
end


text:
struct BoolDistribution
    p::Float64 # probability of true
end

pdf(d::BoolDistribution, s::Bool) = s ? d.p : 1.0-d.p
rand(rng::AbstractRNG, d::BoolDistribution) = rand(rng) <= d.p
iterator(d::BoolDistribution) = [true, false]
Base.:(==)(d1::BoolDistribution, d2::BoolDistribution) = d1.p == d2.p
Base.hash(d::BoolDistribution, u::UInt64=UInt64(0)) = hash(d.p, u)
Base.length(d::BoolDistribution) = 2

@with_kw struct CryingBaby
    r_hungry::Float64 = -10.0
    r_feed::Float64 = -5.0
    r_sing::Float64 = -0.5
    p_become_hungry::Float64 = 0.1
    p_cry_when_hungry::Float64 = 0.8
    p_cry_when_not_hungry::Float64 = 0.1
    p_cry_when_hungry_in_sing::Float64 = 0.9
    γ::Float64 = 0.9
end

# CryingBaby = CryingBaby(-10.0, -5.0, -0.5, 0.1, 0.8, 0.1, 0.9, 0.9)

SATED = 1
HUNGRY = 2
FEED = 1
IGNORE = 2
SING = 3
CRYING = true
QUIET = false

CRYING_BABY_ACTION_COLORS = Dict(
    FEED => "pastelBlue",
    IGNORE => "pastelGreen",
    SING => "pastelRed"
)
CRYING_BABY_ACTION_NAMES = Dict(
    FEED => "feed",
    IGNORE => "ignore",
    SING => "sing",
)

n_states(::CryingBaby) = 2
n_actions(::CryingBaby) = 3
n_observations(::CryingBaby) = 2
discount(pomdp::CryingBaby) = pomdp.γ

ordered_states(::CryingBaby) = [SATED, HUNGRY]
ordered_actions(::CryingBaby) = [FEED, IGNORE, SING]
ordered_observations(::CryingBaby) = [CRYING,QUIET]

two_state_categorical(p1::Float64) = Categorical([p1,1.0 - p1])
function transition(pomdp::CryingBaby, s::Int, a::Int)
    if a == FEED
        return two_state_categorical(1.0)
    else
        if s == HUNGRY
            return two_state_categorical(0.0)
        else
            # Did not feed when not hungry
            return two_state_categorical(1.0-pomdp.p_become_hungry)
        end
    end
end

function observation(pomdp::CryingBaby, a::Int, s′::Int)
    if a == SING
        if s′ == HUNGRY
            return BoolDistribution(pomdp.p_cry_when_hungry_in_sing)
        else
            return BoolDistribution(0.0)
        end
    else # FEED or IGNORE
        if s′ == HUNGRY
            return BoolDistribution(pomdp.p_cry_when_hungry)
        else
            return BoolDistribution(pomdp.p_cry_when_not_hungry)
        end
    end
end

function reward(pomdp::CryingBaby, s::Int, a::Int)
    r = 0.0
    if s == HUNGRY
        r += pomdp.r_hungry
    end
    if a == FEED
        r += pomdp.r_feed
    elseif a == SING
        r += pomdp.r_sing
    end
    return r
end

reward(pomdp::CryingBaby, b::Vector{Float64}, a::Int) = sum(reward(pomdp,s,a)*b[s] for s in ordered_states(pomdp))

function DiscretePOMDP(pomdp::CryingBaby; γ::Float64=pomdp.γ)
    nS = n_states(pomdp)
    nA = n_actions(pomdp)
    nO = n_observations(pomdp)

    T = zeros(nS, nA, nS)
    R = Array{Float64}(undef, nS, nA)
    O = Array{Float64}(undef, nA, nS, nO)

    s_s = 1
    s_h = 2

    a_f = 1
    a_i = 2
    a_s = 3

    o_c = 1
    o_q = 2

    T[s_s, a_f, :] = [1.0, 0.0]
    T[s_s, a_i, :] = [1.0-pomdp.p_become_hungry, pomdp.p_become_hungry]
    T[s_s, a_s, :] = [1.0-pomdp.p_become_hungry, pomdp.p_become_hungry]
    T[s_h, a_f, :] = [1.0, 0.0]
    T[s_h, a_i, :] = [0.0, 1.0]
    T[s_h, a_s, :] = [0.0, 1.0]

    R[s_s, a_f] = reward(pomdp, s_s, a_f)
    R[s_s, a_i] = reward(pomdp, s_s, a_i)
    R[s_s, a_s] = reward(pomdp, s_s, a_s)
    R[s_h, a_f] = reward(pomdp, s_h, a_f)
    R[s_h, a_i] = reward(pomdp, s_h, a_i)
    R[s_h, a_s] = reward(pomdp, s_h, a_s)

    O[a_f, s_s, :] = [observation(pomdp, a_f, s_s).p, 1 - observation(pomdp, a_f, s_s).p]
    O[a_f, s_h, :] = [observation(pomdp, a_f, s_h).p, 1 - observation(pomdp, a_f, s_h).p]
    O[a_i, s_s, :] = [observation(pomdp, a_i, s_s).p, 1 - observation(pomdp, a_i, s_s).p]
    O[a_i, s_h, :] = [observation(pomdp, a_i, s_h).p, 1 - observation(pomdp, a_i, s_h).p]
    O[a_s, s_s, :] = [observation(pomdp, a_s, s_s).p, 1 - observation(pomdp, a_s, s_s).p]
    O[a_s, s_h, :] = [observation(pomdp, a_s, s_h).p, 1 - observation(pomdp, a_s, s_h).p]

    return DiscretePOMDP(T, R, O, γ)
end

function POMDP(pomdp::CryingBaby; γ::Float64=pomdp.γ)
    disc_pomdp = DiscretePOMDP(pomdp)
    return POMDP(disc_pomdp)
end


text:
module TestExamples3DAdvection

using Test
using Trixi

include("test_trixi.jl")

# pathof(Trixi) returns /path/to/Trixi/src/Trixi.jl, dirname gives the parent directory
EXAMPLES_DIR = joinpath(pathof(Trixi) |> dirname |> dirname, "examples", "3d")

@testset "Linear scalar advection" begin
  @trixi_testset "elixir_advection_basic.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_advection_basic.jl"),
      l2   = [0.00015975755208652597],
      linf = [0.0015038732976652147])
  end

  @trixi_testset "elixir_advection_restart.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_advection_restart.jl"),
      l2   = [0.0001780001287314664],
      linf = [0.0014520752637396939])
  end

  @trixi_testset "elixir_advection_extended.jl with initial_condition_sin" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_advection_extended.jl"),
      l2   = [0.002727293067654415],
      linf = [0.024833049753677727],
      initial_condition=Trixi.initial_condition_sin)
  end

  @trixi_testset "elixir_advection_extended.jl with initial_condition_constant" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_advection_extended.jl"),
      l2   = [7.033186506921888e-16],
      linf = [2.6645352591003757e-15],
      initial_condition=initial_condition_constant)
  end

  @trixi_testset "elixir_advection_extended.jl with initial_condition_linear_z and periodicity=false" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_advection_extended.jl"),
      l2   = [5.887699520794518e-16],
      linf = [6.217248937900877e-15],
      initial_condition=Trixi.initial_condition_linear_z,
      boundary_conditions=Trixi.boundary_condition_linear_z, periodicity=false)
  end

  @trixi_testset "elixir_advection_mortar.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_advection_mortar.jl"),
      l2   = [0.0018461529502663268],
      linf = [0.01785420966285467])
  end

  @trixi_testset "elixir_advection_amr.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_advection_amr.jl"),
      l2   = [9.773852895157622e-6],
      linf = [0.0005853874124926162])
  end
end

end # module


text:
using DirichletProcess
using Test

@testset "DirichletProcess.jl" begin
    # Write your tests here.
end


text:
using ITensors, Test

@testset "ITensors.jl" begin
    @testset "$filename" for filename in (
        "test_tagset.jl",
        "test_smallstring.jl",
        "test_index.jl",
        "test_indexset.jl",
        "test_itensor_dense.jl",
        "test_itensor_diag.jl",
        "test_contract.jl",
        "test_combiner.jl",
        "test_trg.jl",
        "test_ctmrg.jl",
        "test_iterativesolvers.jl",
        "test_dmrg.jl",
        "test_tag_types.jl",
        "test_phys_site_types.jl",
        "test_decomp.jl",
        "test_lattices.jl",
        "test_mps.jl",
        "test_mpo.jl",
        "test_autompo.jl",
        "test_svd.jl",
        "test_qn.jl",
        "test_readme.jl",
    )
      println("Running $filename")
      include(filename)
    end
end


text:
import LinearAlgebra: det, logdet, LowerTriangular, UpperTriangular

const ∇ScalarLT = LowerTriangular{<:∇Scalar}
const ∇ScalarUT = UpperTriangular{<:∇Scalar}

for (ctor, T) in zip([:LowerTriangular, :UpperTriangular], [:∇ScalarLT, :∇ScalarUT])

    @eval @explicit_intercepts $ctor Tuple{∇AbstractMatrix}
    @eval ∇(::Type{$ctor}, ::Type{Arg{1}}, p, Y::$T, Ȳ::$T, X::∇AbstractMatrix) = Matrix(Ȳ)
    @eval ∇(
        X̄::∇AbstractMatrix,
        ::Type{$ctor},
        ::Type{Arg{1}},
        p,
        Y::$T,
        Ȳ::$T,
        X::∇AbstractMatrix,
    ) = broadcast!(+, X̄, X̄, Ȳ)

    @eval @explicit_intercepts det Tuple{$T}
    @eval ∇(::typeof(det), ::Type{Arg{1}}, p, y::∇Scalar, ȳ::∇Scalar, X::$T) =
        Diagonal(ȳ .* y ./ view(X, diagind(X)))

    # Optimisation for in-place updates.
    @eval function ∇(
        X̄::$T,
        ::typeof(det),
        ::Type{Arg{1}},
        p,
        y::∇Scalar,
        ȳ::∇Scalar,
        X::$T,
    )
        X̄_diag = view(X̄, diagind(X̄))
        broadcast!((x̄, x, y, ȳ)->x̄ + ȳ * y / x,
                   X̄_diag, X̄_diag, view(X, diagind(X)), y, ȳ)
        return X̄
    end

    # Optimisation for in-place updates to `Diagonal` sensitivity cache.
    @eval function ∇(
        X̄::Diagonal,
        ::typeof(det),
        ::Type{Arg{1}},
        p,
        y::∇Scalar,
        ȳ::∇Scalar,
        X::$T,
    )
        X̄.diag .+= ȳ .* y ./ view(X, diagind(X))
        return X̄
    end

    @eval @explicit_intercepts logdet Tuple{$T}
    @eval ∇(::typeof(logdet), ::Type{Arg{1}}, p, y::∇Scalar, ȳ::∇Scalar, X::$T) =
        Diagonal(ȳ ./ view(X, diagind(X)))

    # Optimisation for in-place updates.
    @eval function ∇(
        X̄::∇Array,
        ::typeof(logdet),
        ::Type{Arg{1}},
        p,
        y::∇Scalar,
        ȳ::∇Scalar,
        X::$T
    )
        X̄_diag = view(X̄, diagind(X̄))
        broadcast!((x̄, x, ȳ)->x̄ + ȳ / x, X̄_diag, X̄_diag, view(X, diagind(X)), ȳ)
        return X̄
    end

    # Optimisation for in-place updates to `Diagonal` sensitivity cache.
    @eval function ∇(
        X̄::Diagonal,
        ::typeof(logdet),
        ::Type{Arg{1}},
        p,
        y::∇Scalar,
        ȳ::∇Scalar,
        X::$T,
    )
        X̄.diag .+= ȳ ./ view(X, diagind(X))
        return X̄
    end
end


text:
struct Procrustes
    R
    t
    function Procrustes(X, Y)
        # https://en.wikipedia.org/wiki/Orthogonal_Procrustes_problem
        mX = mean(X, dims = 1)
        mY = mean(Y, dims = 1)
        # Xb, Yb = X .- mX, Y .- mY
        s = svd((X .- mX)' * (Y .- mY))
        R = s.U * s.V'
        new(R, mY - mX * R)
    end
end

function align(X::AbstractMatrix{T}, anchors::AbstractMatrix{T}) where T
    p = Procrustes(X, anchors)
    return X * p.R .+ p.t
end

function align!(X::AbstractMatrix{S}, anchors::AbstractMatrix{T}) where {S,T}
    p = Procrustes(X, anchors)
    X .= X * p.R .+ p.t
    return nothing
end
align!(X, anchors::Nothing) = X

function align(X::AbstractMatrix{T}, anchors::NamedTuple{(:idx, :pos)}) where T
    xs = X[anchors.idx, :]
    ys = anchors.pos
    p = Procrustes(xs, ys)
    return X * p.R .+ p.t
end

function align!(X::AbstractMatrix{T}, anchors::NamedTuple{(:idx, :pos)}) where T
    xs = X[anchors.idx, :]
    ys = anchors.pos
    p = Procrustes(xs, ys)
    X .= X * p.R .+ p.t
end

transform(X, p::Procrustes) = X * p.R .+ p.t


text:
## Exercise 5-1
## As an exercise, draw a stack diagram for printn called with s = "Hello" and n = 2. Then write a function called do_n that takes a function object and a number, n, as arguments, and that calls the given function n times.
println("Ans: ")

println(" n --> 2")
println(" n --> 1")
println(" n --> 0")

function printhello()
    println("Hello")
end

function do_n(func, n)
    if n <= 0
        return
    end

    func()
    do_n(func, n - 1)
end

do_n(printhello, 3)

println("End.")


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

## promote()

@test promote_type(Some{Int}, Some{Float64}) == Some
@test promote_type(Some{Int}, Some{Real}) == Some{Real}
@test promote_type(Some{Int}, Nothing) == Union{Some{Int},Nothing}

## convert()

# These conversions must fail to prevent ambiguities
# when a value to wrap is already a Some or a Nothing
@test_throws MethodError convert(Some, 1)
@test_throws MethodError convert(Union{Some, Nothing}, 1)
@test_throws MethodError convert(Some{Int}, 1)
@test_throws MethodError convert(Union{Some{Int}, Nothing}, 1)

@test convert(Some, Some(1)) === convert(Union{Some, Nothing}, Some(1)) === Some(1)
@test convert(Some{Int}, Some(1)) === convert(Union{Some{Int}, Nothing}, Some(1)) === Some(1)
@test convert(Some{Int}, Some(1.0)) === Some(1)
@test convert(Union{Some{Int}, Nothing}, Some(1.0)) === Some(1)

@test_throws MethodError convert(Some, nothing)
@test_throws MethodError convert(Some{Int}, nothing)

@test convert(Some, Some(nothing)) === Some(nothing)
@test convert(Some{Nothing}, Some(nothing)) === Some(nothing)

@test convert(Union{Some, Nothing}, nothing) === nothing
@test convert(Union{Some{Int}, Nothing}, nothing) === nothing

@test convert(Union{Int, Nothing}, nothing) === nothing
@test convert(Union{Int, Nothing}, 1) === 1
@test convert(Union{Int, Nothing}, 1.0) === 1
@test convert(Nothing, nothing) === nothing
@test_throws MethodError convert(Nothing, 1)

## show()

@test sprint(show, Some(1)) == "Some(1)"
@test sprint(show, Some(Some(1))) == "Some(Some(1))"
@test repr([Some(1)]) == "Some{$Int}[1]"
@test repr([Some(Some(1))]) == "Some{Some{$Int}}[Some(1)]"

##  == and isequal nothing

@test Some(1) != nothing
@test Some(nothing) != nothing
@test !isequal(Some(1), nothing)
@test !isequal(Some(nothing), nothing)

@testset "something" begin
    @test_throws ArgumentError something()
    @test something(1) === 1
    @test something(missing) === missing
    @test_throws ArgumentError something(nothing)
    @test something(nothing, 1) === 1
    @test something(1, nothing) === 1
    @test_throws ArgumentError something(nothing, nothing)
    @test something(nothing, 1, 2) === 1
    @test something(1, nothing, 2) === 1
    @test something(nothing, nothing, 2) === 2

    @test something(Some(1)) === 1
    @test something(Some(nothing)) === nothing
    @test something(Some(missing)) === missing
    @test something(Some(1), 0) === 1
    @test something(Some(nothing), 0) === nothing
    @test something(nothing, Some(nothing)) === nothing
    @test something(Some(1), nothing) === 1
    @test something(nothing, Some(1)) === 1
    @test something(nothing, Some(1), nothing) === 1
    @test something(nothing, Some(1), Some(2)) === 1
    @test something(Some(1), nothing, Some(2)) === 1

    @test something(nothing, missing) === missing
    @test something(missing, nothing) === missing
    @test something(nothing, missing, nothing) === missing
    @test something(missing, nothing, missing) === missing
end

# issue #26927
a = [missing, nothing, Some(nothing), Some(missing)]
@test a isa Vector{Union{Missing, Nothing, Some}}
@test a[1] === missing && a[2] === nothing && a[3] === Some(nothing) && a[4] === Some(missing)
b = [ "replacement", "replacement", nothing, missing ]
@test b isa Vector{Union{Missing, Nothing, String}}
# the original operation from the issue, though it was not the source of the problem
@test all(coalesce.(a, "replacement") .=== ["replacement", nothing, Some(nothing), Some(missing)])
@test all(something.(a, "replacement") .=== [missing, "replacement", nothing, missing])

# notnothing()

using Base: notnothing
@test notnothing(1) === 1
@test_throws ArgumentError notnothing(nothing)

# isnothing()
@test !isnothing(1)
@test isnothing(nothing)

max_stars_repo_name: milesfrain/julia

max_issues_repo_name: AsForMe/julia

max_forks_repo_name: AsForMe/julia


text:
using Distributions, Test, Random, LinearAlgebra, FillArrays
using Distributions: Product

@testset "Testing normal product distributions" begin
    Random.seed!(123456)
    N = 11
    # Construct independent distributions and `Product` distribution from these.
    μ = randn(N)
    ds = Normal.(μ, 1.0)
    x = rand.(ds)
    d_product = product_distribution(ds)
    @test d_product isa MvNormal
    # Check that methods for `Product` are consistent.
    @test length(d_product) == length(ds)
    @test eltype(d_product) === eltype(ds[1])
    @test logpdf(d_product, x) ≈ sum(logpdf.(ds, x))
    @test mean(d_product) == mean.(ds)
    @test var(d_product) == var.(ds)
    @test cov(d_product) == Diagonal(var.(ds))
    @test entropy(d_product) ≈ sum(entropy.(ds))

    y = rand(d_product)
    @test y isa typeof(x)
    @test length(y) == N
end

@testset "Testing generic product distributions" begin
    Random.seed!(123456)
    N = 11
    # Construct independent distributions and `Product` distribution from these.
    ubound = rand(N)
    ds = Uniform.(-ubound, ubound)
    x = rand.(ds)
    d_product = product_distribution(ds)
    @test d_product isa Product
    # Check that methods for `Product` are consistent.
    @test length(d_product) == length(ds)
    @test eltype(d_product) === eltype(ds[1])
    @test logpdf(d_product, x) ≈ sum(logpdf.(ds, x))
    @test mean(d_product) == mean.(ds)
    @test var(d_product) == var.(ds)
    @test cov(d_product) == Diagonal(var.(ds))
    @test entropy(d_product) == sum(entropy.(ds))
    @test insupport(d_product, ubound) == true
    @test insupport(d_product, ubound .+ 1) == false
    @test minimum(d_product) == -ubound
    @test maximum(d_product) == ubound
    @test extrema(d_product) == (-ubound, ubound)
    @test isless(extrema(d_product)...)

    y = rand(d_product)
    @test y isa typeof(x)
    @test length(y) == N
end

@testset "Testing discrete non-parametric product distribution" begin
    Random.seed!(123456)
    N = 11

    for a in ([0, 1], [-0.5, 0.5])
        # Construct independent distributions and `Product` distribution from these.
        support = fill(a, N)
        ds = DiscreteNonParametric.(support, Ref([0.5, 0.5]))
        x = rand.(ds)
        d_product = product_distribution(ds)
        @test d_product isa Product
        # Check that methods for `Product` are consistent.
        @test length(d_product) == length(ds)
        @test eltype(d_product) === eltype(ds[1])
        @test logpdf(d_product, x) ≈ sum(logpdf.(ds, x))
        @test mean(d_product) == mean.(ds)
        @test var(d_product) == var.(ds)
        @test cov(d_product) == Diagonal(var.(ds))
        @test entropy(d_product) == sum(entropy.(ds))
        @test insupport(d_product, fill(a[2], N)) == true
        @test insupport(d_product, fill(a[2] + 1, N)) == false

        y = rand(d_product)
        @test y isa typeof(x)
        @test length(y) == N
    end
end

@testset "Testing iid product distributions" begin
    Random.seed!(123456)
    N = 11
    d = Product(Fill(Laplace(0.0, 2.3), N))
    @test N == length(unique(rand(d)));
    @test mean(d) === Fill(0.0, N)
    @test cov(d) === Diagonal(Fill(var(Laplace(0.0, 2.3)), N))
end


text:
struct DictDiff{T1, KT1, T2, KT2} <: DeepDiff
    before::T1
    after::T2
    removed::Set{KT1}
    added::Set{KT2}
    changed::Dict{KT1, DeepDiff}
    unchanged::Set{KT1}
end

before(diff::DictDiff) = diff.before
after(diff::DictDiff) = diff.after
removed(diff::DictDiff) = diff.removed
added(diff::DictDiff) = diff.added
changed(diff::DictDiff) = diff.changed

Base.:(==)(lhs::DictDiff, rhs::DictDiff) = fieldequal(lhs, rhs)

function deepdiff(X::AbstractDict, Y::AbstractDict)
    xkeys = Set(keys(X))
    ykeys = Set(keys(Y))
    bothkeys = intersect(xkeys, ykeys)

    removed = setdiff(xkeys, ykeys)
    added = setdiff(ykeys, xkeys)
    unchanged = Set{eltype(bothkeys)}()
    changed = Dict{eltype(bothkeys), DeepDiff}()

    for key in bothkeys
        if X[key] != Y[key]
            changed[key] = deepdiff(X[key], Y[key])
        else
            push!(unchanged, key)
        end
    end

    DictDiff(X, Y, removed, added, changed, unchanged)
end

Base.show(io::IO, diff::DictDiff) = diffprint(io, diff, 0)

# indentation space
const inspace = "    "

function diffprint(io, d::DictDiff, indent=0)
    bef = before(d)
    aft = after(d)

    println(io, "Dict(")
    for k in d.unchanged
        # extra space to account for added linemarker
        print(io, " ", inspace ^ (indent+1))
        prettyprint(io, Pair(k, bef[k]), " ", indent+1)
        println(io, ",")
    end
    Base.with_output_color(:red, io) do io
        for k in removed(d)
            print(io, "-", inspace ^ (indent+1))
            prettyprint(io, Pair(k, bef[k]), "-", indent+1)
            println(io, ",")
        end
    end
    for (k, v) in changed(d)
        if isa(v, SimpleDiff)
            # if we have a key pointing to a SimpleDiff, then we don't know how to
            # deconstruct the value, so instead we print it like a removed and added key
            Base.with_output_color(:red, io) do io
                print(io, "-", inspace ^ (indent+1))
                prettyprint(io, Pair(k, before(v)), "-", indent+1)
                println(io, ",")
            end
            Base.with_output_color(:green, io) do io
                print(io, "+", inspace ^ (indent+1))
                prettyprint(io, Pair(k, after(v)), "+", indent+1)
                println(io, ",")
            end
        else
            # extra space to account for added linemarker
            print(io, " ", inspace ^ (indent+1))
            prettyprint(io, Pair(k, v), " ", indent+1)
            println(io, ",")
        end
    end
    Base.with_output_color(:green, io) do io
        for k in added(d)
            print(io, "+", inspace ^ (indent+1))
            prettyprint(io, Pair(k, aft[k]), "+", indent+1)
            println(io, ",")
        end
    end
    # don't print the leading space if we're at the top-level
    print(io, indent == 0 ? "" : " ")
    print(io, inspace ^ indent, ")")
end

function prettyprint(io, d::AbstractDict, linemarker, indent)
    println(io, "Dict(")
    for p in d
        print(io, linemarker, inspace ^ (indent+1))
        prettyprint(io, p, linemarker, indent+1)
        println(io, ",")
    end
    print(io, linemarker, inspace ^ indent, ")")
end

function prettyprint(io, p::Pair, linemarker, indent)
    prettyprint(io, p[1], linemarker, indent)
    print(io, " => ")
    prettyprint(io, p[2], linemarker, indent)
end

function prettyprint(io, p::Pair{<:Any, <:DictDiff}, linemarker, indent)
    prettyprint(io, p[1], linemarker, indent)
    print(io, " => ")
    diffprint(io, p[2], indent)
end

prettyprint(io, x, linemarker, indent) = show(io, x)


text:
using Base: showarg

function summary_build(A)
    # This is just a copy of Base.summary, but copying it prevents a
    # StackOverflow
    io = IOBuffer()
    summary(io, A)
    String(take!(io))
end
function summary_build(A, cthresh)
    Base.depwarn("summary_build(A, cthresh) is deprecated (and cthresh isn't used), just use summary(A)", :summary_build)
    summary_build(A)
end

# Not perfect, but better to have this than not
Base.showarg(io::IO, x) = showarg(io, x, false)

function type_complexity(x)
    Base.depwarn("type_complexity is deprecated, please use 1", :type_complexity)
    1
end


text: println("I'm hungry!")


text:
"""
    σ(x) = 1 / (1 + exp(-x))

Classic [sigmoid](https://en.wikipedia.org/wiki/Sigmoid_function) activation
function.
"""
σ(x) = one(x) / (one(x) + exp(-x))

const sigmoid = σ

# ForwardDiff numerical stability hack
σ_stable(x) = ifelse(x < -80, zero(x), one(x) / (one(x) + exp(-x)))

σ(x::Float32) = σ_stable(x)

@require ForwardDiff begin
  import ForwardDiff: Dual
  σ(x::Dual{T,Float32}) where T = σ_stable(x)
end

"""
    logσ(x)

Return `log(σ(x))` which is computed in a numerically stable way.

    julia> logσ(0.)
    -0.6931471805599453
    julia> logσ.([-100, -10, 100.])
    3-element Array{Float64,1}:
     -100.0
      -10.0
       -0.0
"""
function logσ(x)
  max_v = max(zero(x), -x)
  z = exp(-max_v) + exp(-x-max_v)
  -(max_v + log(z))
end

const logsigmoid = logσ

"""
    relu(x) = max(0, x)

[Rectified Linear Unit](https://en.wikipedia.org/wiki/Rectifier_(neural_networks))
activation function.
"""
relu(x) = max(zero(x), x)

"""
    leakyrelu(x) = max(0.01x, x)

Leaky [Rectified Linear Unit](https://en.wikipedia.org/wiki/Rectifier_(neural_networks))
activation function.
You can also specify the coefficient explicitly, e.g. `leakyrelu(x, 0.01)`.
"""
leakyrelu(x, a = oftype(x/1, 0.01)) = max(a*x, x/1)

"""
    elu(x, α = 1) =
      x > 0 ? x : α * (exp(x) - 1)

Exponential Linear Unit activation function.
See [Fast and Accurate Deep Network Learning by Exponential Linear Units](https://arxiv.org/abs/1511.07289).
You can also specify the coefficient explicitly, e.g. `elu(x, 1)`.
"""
elu(x, α = one(x)) = ifelse(x ≥ 0, x/1, α * (exp(x) - one(x)))

"""
    swish(x) = x * σ(x)

Self-gated actvation function.
See [Swish: a Self-Gated Activation Function](https://arxiv.org/pdf/1710.05941.pdf).
"""
swish(x) = x * σ(x)

"""
    selu(x) = λ * (x ≥ 0 ? x : α * (exp(x) - 1))

    λ ≈ 1.0507
    α ≈ 1.6733

Scaled exponential linear units.
See [Self-Normalizing Neural Networks](https://arxiv.org/pdf/1706.02515.pdf).
"""
function selu(x)
  λ = oftype(x/1, 1.0507009873554804934193349852946)
  α = oftype(x/1, 1.6732632423543772848170429916717)
  λ * ifelse(x > 0, x/1, α * (exp(x) - 1))
end

"""
    softsign(x) = x / (1 + |x|)

See [Quadratic Polynomials Learn Better Image Features](http://www.iro.umontreal.ca/~lisa/publications2/index.php/attachments/single/205).
"""
softsign(x) = x / (one(x) + abs(x))

"""
    softplus(x) = log(exp(x) + 1)

See [Deep Sparse Rectifier Neural Networks](http://proceedings.mlr.press/v15/glorot11a/glorot11a.pdf).
"""
softplus(x) = log1p(exp(x))


text:
using Mixin
using Test

@testset "Mixin.jl" begin
    get_rettyp(f, argtypes) = last(first(code_typed(f, argtypes)))

    # basic
    @interface struct Ty1
        con::Int
        abs::Union{Symbol,String}
        any
   end

    @mixin struct Ty11 <: Ty1
        con::Int
        abs::Union{Symbol,String}
        any
    end

    @mixin struct Ty12 <: Ty1
        con::Int
        abs::Symbol
        any
        additional::Symbol
        function Ty12(con, args...)
            con ≤ 0 && error("con should be positive for some reason")
            new(con, args...)
        end
    end

    @test get_rettyp(ary->first(ary).con, (Vector{Ty1},)) == Int
    @test get_rettyp(ary->first(ary).abs, (Vector{Ty1},)) == Union{Symbol,String}

    # FIXME: we want to check against AssertionError if possible

    # lack of required field
    @test_throws LoadError macroexpand(@__MODULE__, :(
    @mixin struct Ty13 <: Ty1
        con::Int
        any
    end
    ))

    # invalid typing
    @test_throws LoadError macroexpand(@__MODULE__, :(
    @mixin struct Ty14 <: Ty1
        con::Int
        abs::Int
        any
    end
    ))
end


text:
using KNITRO, Compat.Test

#    min  9 - 8x1 - 6x2 - 4x3
#         + 2(x1^2) + 2(x2^2) + (x3^2) + 2(x1*x2) + 2(x1*x3)
#    subject to  c[0]:  x1 + x2 + 2x3 <= 3
#                x1 >= 0
#                x2 >= 0
#                x3 >= 0
#    initpt (0.5, 0.5, 0.5)
#
#    Solution is x1=4/3, x2=7/9, x3=4/9, lambda=2/9  (f* = 1/9)
#
#  The problem comes from Hock and Schittkowski, HS35.

function eval_f(x::Vector{Float64})
    linear_terms = 9.0 - 8.0*x[1] - 6.0*x[2] - 4.0*x[3]
    quad_terms = 2.0*x[1]^2 + 2.0*x[2]^2 + x[3]^2 + 2.0*x[1]*x[2] + 2.0*x[1]*x[3]
    return linear_terms + quad_terms
end

function eval_g(x::Vector{Float64}, cons::Vector{Float64})
    cons[1] = x[1] + x[2] + 2.0*x[3]
end

function eval_grad_f(x::Vector{Float64}, grad::Vector{Float64})
    grad[1] = -8.0 + 4.0*x[1] + 2.0*x[2] + 2.0*x[3]
    grad[2] = -6.0 + 2.0*x[1] + 4.0*x[2]
    grad[3] = -4.0 + 2.0*x[1]            + 2.0*x[3]
end

function eval_jac_g(x::Vector{Float64}, jac::Vector{Float64})
    jac[1] = 1.0
    jac[2] = 1.0
    jac[3] = 2.0
end

function eval_h(x::Vector{Float64}, lambda::Vector{Float64},
                sigma::Float64, hess::Vector{Float64})
    hess[1] = sigma*4.0
    hess[2] = sigma*2.0
    hess[3] = sigma*2.0
    hess[4] = sigma*4.0
    hess[5] = sigma*2.0
end

function eval_hv(x::Vector{Float64}, lambda::Vector{Float64},
                 sigma::Float64, hv::Vector{Float64})
    hv[1] = sigma*4.0*hv[1] + sigma*2.0*hv[2] + sigma*2.0*hv[3]
    hv[2] = sigma*2.0*hv[1] + sigma*4.0*hv[2]
    hv[3] = sigma*2.0*hv[1]                   + sigma*2.0*hv[3]
end

objGoal = KTR_OBJGOAL_MINIMIZE
objType = KTR_OBJTYPE_QUADRATIC

n = 3
x_L = zeros(n)
x_U = [KTR_INFBOUND,KTR_INFBOUND,KTR_INFBOUND]

m = 1
c_Type = [KTR_CONTYPE_LINEAR]
c_L = [-KTR_INFBOUND]
c_U = [3.0]

jac_con = Int32[0,0,0]
jac_var = Int32[0,1,2]
hess_row = Int32[0,0,0,1,2]
hess_col = Int32[0,1,2,1,2]

x       = [0.5,0.5,0.5]
lambda  = zeros(n+m)
obj     = [0.0]

kp = createProblem()
@test applicationReturnStatus(kp) == :Uninitialized
loadOptionsFile(kp,joinpath(dirname(@__FILE__),"knitro.opt"))
@test applicationReturnStatus(kp) == :Uninitialized

# --- test getters and setters ---
hessopt = zeros(Int32, 1)
getOption(kp, "hessopt", hessopt)
@test hessopt[1] == 1
setOption(kp, "hessopt", 0)
getOption(kp, "hessopt", hessopt)
@test hessopt[1] == 0
setOption(kp, "hessopt", 1)
getOption(kp, "hessopt", hessopt)
@test hessopt[1] == 1
# ---------------------------------

initializeProblem(kp, objGoal, objType, x_L, x_U, c_Type, c_L, c_U,
                  jac_var, jac_con, hess_row, hess_col)
@test applicationReturnStatus(kp) == :Initialized
setCallbacks(kp, eval_f, eval_g, eval_grad_f, eval_jac_g, eval_h, eval_hv)
@test applicationReturnStatus(kp) == :Initialized
solveProblem(kp)

# --- test optimal solutions ---
@testset "Test optimal solutions" begin
    @test applicationReturnStatus(kp) == :Optimal
    @test isapprox(kp.x, [1.3333333, 0.7777777, 0.4444444], atol=1e-5)
    @test isapprox(kp.obj_val[1], 0.1111111, atol=1e-5)
end

@testset "Test getters" begin
    @test get_number_FC_evals(kp) >= 0
    @test get_number_GA_evals(kp) >= 0
    @test get_number_H_evals(kp) >= 0
    @test get_number_HV_evals(kp) >= 0
    @test get_number_iters(kp) >= 0
    @test get_number_cg_iters(kp) >= 0
    @test isapprox(get_abs_feas_error(kp), 0.0, atol=1e-5)
    @test isapprox(get_rel_feas_error(kp), 0.0, atol=1e-5)
    @test isapprox(get_abs_opt_error(kp), 0.0, atol=1e-5)
    @test isapprox(get_rel_opt_error(kp), 0.0, atol=1e-5)
end

# start at a different starting point
x = [2.0,2.0,2.0]
lambda = ones(n+m)
restartProblem(kp, x, lambda)
@test applicationReturnStatus(kp) == :Initialized
solveProblem(kp)
@testset "Test optimal solutions (at a different starting point)" begin
    @test applicationReturnStatus(kp) == :Optimal
    @test isapprox(kp.x,  [1.3333333, 0.7777777, 0.4444444], atol=1e-5)
    @test isapprox(kp.obj_val[1], 0.1111111, atol=1e-5)
end

freeProblem(kp)