text:
"""
    subsampleshift(waveform, quantity, [algorithm=BSpline(Cubic(Interpolations.Flat(OnGrid())))])

Return `waveform` shifted by `quantity` using `algorithm` for interpolation.

# Examples

```jldoctest
julia> let f(x) = x^3 + 2x^2 + 3x + 4, x = range(0., 1., 256)
       isapprox(f.(x), subsampleshift(subsampleshift(f.(x), 0.5), -0.5), atol=1e-2)
       end
true
```
"""
function subsampleshift(waveform::AbstractVector{<:AbstractFloat}, quantity; algorithm=BSpline(Cubic(Interpolations.Flat(OnGrid()))))
	if abs(quantity) > 1.0
		throw(ArgumentError("Shift should be less than 1.0, got $quantity"))
	else
		xs = axes(waveform, 1)
		itp = interpolate(waveform, algorithm)
		ext = extrapolate(itp, Interpolations.Flat())
		@. ext(xs + quantity)
	end
end

"""
    findmax_window(crosscorrelation, index, tolerance)

Return a tuple containing the maximum of `crosscorrelation`, in a window
centered around `index` and radius equal to `tolerace`, and its index.

# Examples

```jldoctest
julia> findmax_window.([abs.(sin.(0:0.01pi:4pi)) for _ = 1:4], [49, 153, 249, 353], 5)
4-element Vector{Tuple{Float64, Int64}}:
 (1.0, 51)
 (1.0, 151)
 (1.0, 251)
 (1.0, 351)
```
"""
function findmax_window(crosscorrelation, index, tolerance)
    if isempty(crosscorrelation)
        throw(ArgumentError("Cross-correlation series must be non-empty."))
    elseif !(index in axes(crosscorrelation, 1))
        throw(ArgumentError("Index must be a valid index of cross-correlation series."))
    end
    start = max(index  - tolerance, firstindex(crosscorrelation))
    stop = min(index  + tolerance, lastindex(crosscorrelation))
    if start == firstindex(crosscorrelation) && stop == lastindex(crosscorrelation)
        findmax(crosscorrelation)
    else
        cc_max, indx_max = findmax(view(crosscorrelation, start:stop))
        cc_max, start + indx_max - 1
    end
end

"""
    estimatetoa(trace, waveform, center, left, tolerance)

Return a tuple containing the estimated position, with subsample precision,
and value of maximum cross-correlation between `trace` and `waveform`
in a window around `center` and radius equal to `tolerace`.
The estimation is performed by approximating the cross-correlation with
a parabola near the maximum. This estimate is biased and its accuracy
decreases with decreasing sampling frequencies.
The position of the maximum cross-correlation is the Time Of Arrival (TOA)
in sample units, i.e. the true TOA divided by sampling frequency.

# Examples

```jldoctest
julia> let v = sin.(range(0, pi, 256)), x = [zeros(128); v; zeros(128)], y = TemplateMatching.subsampleshift(v, 0.5)
       toa, cc = estimatetoa(x, y, 128, 5)
       round(toa, digits=1), round(cc, digits=1)
       end
(129.5, 1.0)
```
"""
function estimatetoa(trace, waveform, center, tol)
	a, b = max(firstindex(trace), center - tol), min(lastindex(trace), center + tol + length(waveform) - 1)
	cc = OffsetVector(crosscorrelate(view(trace, a:b), waveform), center - tol - 1)
	cc_max, n_max = findmax_window(cc, center, tol)
	if firstindex(cc) < n_max < lastindex(cc) && cc_max < 1.0 && cc_max ≉ 1.0
		y1, y2, y3 = cc[n_max - 1:n_max + 1]
		delta = (y1 - y3) / 2(y1 - 2y2 + y3)
		n_max + delta, min(1.0, y2 + (y3 - y1) / 4delta)
	else
		Float64(n_max), cc_max
	end
end

function line_element(xt, v)
    x = view(xt, 1:3)
    t = xt[4]
    dot(x, x) - v^2 * t^2
end

residue_rms(xt, sensors_readings_itr, v) = sqrt(mean(ys -> line_element(ys - xt, v)^2, sensors_readings_itr))

locate(sensors_readings_itr, v, guess) = optimize(xt -> residue_rms(xt, sensors_readings_itr, v), guess)


text:
import BasisFunctions: diff
using Calculus: derivative
diff(f::typeof(cos),args...) = derivative(f,args...)
function test_derived_dicts(T)
    b1 = Fourier{T}(11)
    b2 = ChebyshevT{T}(12)

    @testset "Generic derived dictionary" begin
        test_generic_dict_interface(BasisFunctions.ConcreteDerivedDict(b2))
    end

    @testset "Complexified dictionary" begin
        # Something is wrong with BigFloat, the idct isn't right. Consult with Daan.
        b3 = ChebyshevT{Float64}(12)
        # First, check that a complexified dict is made
        b3c = complex(b3)
        @test b3c isa BasisFunctions.ComplexifiedDict
        @test coefficienttype(b3c) == Complex{Float64}
        # and check that it doesn't create a nested ComplexifiedDict
        @test complex(b3c) isa BasisFunctions.ComplexifiedDict{<:ChebyshevT}
        test_generic_dict_interface(b3c)
    end
    @testset "Linear mapped dictionaries" begin
        test_generic_dict_interface(rescale(b1, -T(1), T(2)))
        test_generic_dict_interface(rescale(b2, -T(2), T(3)))
    end

    @testset "A simple subdict" begin
        test_generic_dict_interface(b1[2:6]) end

    @testset "Operated dictionaries" begin
        D1 = DiagonalOperator(b1, rand(length(b1)))
        test_generic_dict_interface(D1*b1)
        D2 = ArrayOperator(rand(length(b1),length(b1)), b1, b1)
        test_generic_dict_interface(D2*b1)
        D3 = differentiation(b2)
        test_generic_dict_interface(D3*b2)
    end

    @testset "Multiple dictionaries" begin
        # Test sets with internal representation as vector and as tuple separately
        test_generic_dict_interface(multidict(b1,rescale(b2, 0, 1)))
        #test_generic_dict_interface(MultiDict((b1,rescale(b2, 0, 1)))) end
        end

    @testset "Weighted dictionaries" begin
        # as well as a regular function
        test_generic_dict_interface(cos * b1)
        # and a 2D example
        # (not just yet, fix 2D for BigFloat first)
#        test_generic_dict_interface( ((x,y) -> cos(x+y)) * tensorproduct(b1, 2) )
    end

    @testset "A multiple and weighted dict combination" begin
        s = rescale(b1, 1/2, 1)
        test_generic_dict_interface(multidict(s,log*s)) end

    @testset "A complicated subdict" begin
        s = rescale(b1, 1/2, 1)
        test_generic_dict_interface(s[1:5]) end

    @testset "Piecewise dictionaries" begin
        # Changed the number of partitions to 7 (from 10) in order to avoid alignment of a
        # grid point in one of the tests with the boundary of a subinterval.
        part = PiecewiseInterval(T(0), T(10), 7)
        pw = PiecewiseDict(b2, part)
        test_generic_dict_interface(pw) end

    # @testset "A tensor product of MultiDict's" begin
    # b = multidict(b1,b2)
    # c = b ⊗ b
    # test_generic_dict_interface(c) end
end

ext: jl


text:
# * System

"""
    System

Describes the combined system of atom/molecule (in terms of ionization
channels) and an external, time-dependent field. The ionization
channels may be coupled through various interactions, which may or may
not affect the photoelectron. `System` only describes the channels,
the external field, and the interactions possible, the actual process
is described by a [`Diagram`](@ref).
"""
struct System{T,IonizationChannels<:AbstractVector{<:IonizationChannel{T}},
              Couplings<:AbstractVector{<:AbstractMatrix{<:AbstractCoupling}},
              VectorPotential,
              Times<:AbstractRange,
              Volkov<:VolkovPhases}
    ionization_channels::IonizationChannels

    couplings::Couplings

    𝐀::VectorPotential
    t::Times
    dt::T

    volkov::Volkov
end

"""
    System(ionization_channels, couplings, F, ndt)

Set up a [`System`](@ref) consisting of multiple
[`IonizationChannel`](@ref)s with possible `couplings` between them,
driven by an external field `F`, sampled at a frequency of `fs`.
"""
function System(ionization_channels::AbstractVector{<:IonizationChannel},
                couplings::AbstractVector,
                F::ElectricFields.AbstractField, fs::Number)
    t = timeaxis(F, fs)
    volkov = VolkovPhases(F, t)

    𝐀 = vector_potential.(F, t)
    System(ionization_channels, couplings, 𝐀, t, step(t), volkov)
end

@doc raw"""
    System(ionization_channels, couplings, Fv, Av, t)

Set up a [`System`](@ref) consisting of multiple
[`IonizationChannel`](@ref)s with possible `couplings` between them,
driven by an external field `Fv` with corresponding vector potential
`Av`, both resolved on the times given by `t`; it is up to the user to
ensure that ``\vec{F} = -\partial_t\vec{A}`` holds.
"""
function System(ionization_channels::AbstractVector{<:IonizationChannel},
                couplings::AbstractVector,
                Fv::AbstractVector, Av::AbstractVector, t::AbstractRange)
    volkov = VolkovPhases(Av, t)

    System(ionization_channels, couplings, Av, t, step(t), volkov)
end

System(ionization_channels::AbstractVector{<:IonizationChannel}, args...;
       couplings=Matrix{AbstractCoupling}[], kwargs...) =
           System(ionization_channels, couplings, args...)

System(ionization_channel::IonizationChannel, args...; kwargs...) =
    System([ionization_channel], args...; kwargs...)

System(Iₚ::Number, args...; kwargs...) =
    System(IonizationChannel(Iₚ, args...), args...; kwargs...)

function Base.show(io::IO, system::System)
    n = length(system.ionization_channels)
    write(io, "$(n)-channel System")
end

function Base.show(io::IO, mime::MIME"text/plain", system::System)
    println(io, system, ":")
    for (i,c) in enumerate(system.ionization_channels)
        println(io, " ", i, ". ", c)
    end
end

canonical_momentum_conservation(system::System, which::Integer) =
    iszero(which) ?
    CanonicalMomentumConservation() :
    canonical_momentum_conservation(first(system.couplings[which]))

kinematic_momentum(k::Number, A::Number) = k + A
kinematic_momentum(k::SVector{3}, A::SVector{3}) = k + A
kinematic_momentum(k::SVector{3}, A::Number) = SVector{3}(k[1], k[2], k[3]+A)

# * Diagrams

"""
    Diagram(path, couplings)

Goldstone diagram describing a strong-field process, i.e. after an
initial photoionization event, the ion and photoelectron are
propagated separately (at the chosen level of approximation, typically
Volkov waves are used for the photoelectrons). By interacting with the
external field, the ion state may change, and through electron
(re)scattering, both photoelectron and ion state may change.

The diagram is specified using a `path` of pairs (in antichronological
order), where each pair designates `(resultant ion channel,
interaction)`, where `interaction` is an index to one of the
`couplings`. The special value `interaction == 0` corresponds to the
initial ionization event, and should appear once at the end of `path`
(i.e. the first event); if it also appears as the first element of
`path` (i.e. the last event), it corresponds to recombination to the
initial state.

# Examples

```jldoctest
julia> couplings = (StrongFieldApproximation.DipoleCoupling,
                    StrongFieldApproximation.CoulombCoupling)
(StrongFieldApproximation.DipoleCoupling, StrongFieldApproximation.CoulombCoupling)

julia> Diagram([(1,0)], couplings)
Goldstone Diagram:
   |0⟩
  ╱   ╲⇜
 1┃   │𝐤

julia> Diagram([(1,0),(1,0)], couplings)
Goldstone Diagram:
   |0⟩
  ╱   ╲⇜
 1┃   │𝐩
  ╲   ╱⇝
   |0⟩

julia> Diagram([(2,1),(1,0)], couplings)
Goldstone Diagram:
   |0⟩
  ╱   ╲⇜
 1┃   │𝐩
 ⇝┃   │
 2┃   │𝐤

julia> Diagram([(2,2),(1,0)], couplings)
Goldstone Diagram:
   |0⟩
  ╱   ╲⇜
 1┃   │𝐩
  ┠┈┈┈┤
 2┃   │𝐤

julia> Diagram([(1,2),(2,2),(1,0)], couplings)
Goldstone Diagram:
   |0⟩
  ╱   ╲⇜
 1┃   │𝐩
  ┠┈┈┈┤
 2┃   │𝐪
  ┠┈┈┈┤
 1┃   │𝐤

julia> Diagram([(3,1),(1,1),(2,2),(1,0)], couplings)
Goldstone Diagram:
   |0⟩
  ╱   ╲⇜
 1┃   │𝐩
  ┠┈┈┈┤
 2┃   │𝐪
 ⇝┃   │
 1┃   │
 ⇝┃   │
 3┃   │𝐤
```
"""
struct Diagram{Couplings}
    path::Vector{Tuple{Int,Int}}
    couplings::Couplings
    function Diagram(path, couplings::Couplings) where Couplings
        if !isempty(path)
            α,which = last(path)
            which == 0 ||
                throw(ArgumentError("Non-empty diagrams must have photoionization as the first interaction"))
        end
        for i = 2:length(path)-1
            path[i][2] == 0 &&
                throw(ArgumentError("Only the first and last interaction in a diagram may have which = 0 (corresponding to ionization/recombination)"))
        end
        if length(path) > 1 && path[1][2] == 0
            path[1][1] == path[2][1] ||
                throw(ArgumentError("Must recombine from the last channel"))
        end
        new{Couplings}(path, couplings)
    end
end

Diagram(path::Tuple{Int,Int}, couplings) =
    Diagram([path], couplings)

"""
    Diagram(path::AbstractVector{Tuple{Int,Int}}, system)

Convenience constructor that sets up the [`Diagram`](@ref) given a
`path` and the couplings present in `system`. If `path` is empty, the
diagram corresponding to photoionization into the first ionization of
channel of `system` will automatically be generated.
"""
function Diagram(path::AbstractVector, system::System)
    if isempty(path)
        length(system.ionization_channels) > 1 &&
            @warn "More than one ionization channel present in system, choosing the first"
        path = [(1,0)]
    end
    Diagram(path, [typeof(first(c)) for c in system.couplings])
end

Diagram(system::System) = Diagram([], system)

for f in [:length, :first, :firstindex, :lastindex, :isempty]
    @eval Base.$f(d::Diagram) = $f(d.path)
end
Base.getindex(d::Diagram, i) = Diagram(d.path[i], d.couplings)

function draw_ionization(io, nd)
    println(io, lpad("|0⟩", nd+4))
    print(io, lpad("╱ ╲", nd+4))
    printstyled(io, "⇜", color=:light_red)
    println(io)
end

function draw_recombination(io, nd)
    print(io, lpad("╲ ╱", nd+4))
    printstyled(io, "⇝", color=:light_red)
    println(io)
    println(io, lpad("|0⟩", nd+4))
end

draw_exciton(io, ion, nd, electron="") =
    println(io, lpad(ion, nd)*"┃   │$(electron)")

draw_coulomb_interaction(io, nd) =
    println(io, repeat(" ",nd)*"┠┈┈┈┤")

function draw_dipole_interaction(io, nd)
    printstyled(io, lpad("⇝",nd), color=:light_red)
    println(io, "┃   │")
end

function Base.show(io::IO, ::MIME"text/plain", d::Diagram)
    println(io, "Goldstone Diagram:")
    if isempty(d)
        println(io, "|0⟩")
        return
    end
    nd = maximum(iw -> length(digits(iw[1])), d.path)+1
    electrons = ["𝐩","𝐪"]
    cur_electron = 1
    for (i,(ion,which)) in Iterators.reverse(enumerate(d.path))
        if which == 0
            if i == length(d.path)
                draw_ionization(io, nd)
            elseif i == 1
                draw_recombination(io, nd)
            end
        else
            c = d.couplings[which]
            if c <: CoulombCoupling
                draw_coulomb_interaction(io, nd)
            elseif c <: DipoleCoupling
                draw_dipole_interaction(io, nd)
            end
        end
        electron = if i == 1
            "𝐤"
        elseif !iszero(which) && d.couplings[which] <: DipoleCoupling
            ""
        else
            ei = mod1(cur_electron,length(electrons))
            e = electrons[ei]*repeat("′", fld1(cur_electron,length(electrons))-1)
            cur_electron += 1
            e
        end

        which == 0 && i == 1 && length(d) > 1 || draw_exciton(io, ion, nd, electron)
    end
end

# * Integrate diagrams

function analyze_diagram(system, diagram)
    α,which = first(diagram)
    ions = Int[]
    # For photoelectron spectra, time 1 is the reference time,
    # typically at which the laser pulse has ended; for direct
    # photoelectron diagrams, time 2 is the time of ionization. For
    # dipoles, time 1 is the time of recombination.
    unique_momenta = Tuple{Int,Int}[(1,2)]
    momenta = [1]
    indeterminate_momenta = Int[]
    ld = length(diagram)
    order = ld

    if ld > 1
        if which == 0
            push!(indeterminate_momenta, 1)
            # Recombination does not increase the order of the diagram, since
            # it only amounts to projecting the wavefunction on ⟨0|𝐫.
            order -= 1
        end
    else
        push!(ions, α)
    end

    i = 2
    for (α,which) in diagram.path[(iszero(which) ? 2 : 1):end]
        push!(ions, α)
        a,b = unique_momenta[end]
        unique_momenta[end] = (a,i)
        if !(iszero(which) || canonical_momentum_conservation(system, which) == CanonicalMomentumConservation())
            push!(unique_momenta, (i,i+1))
            push!(indeterminate_momenta, length(unique_momenta))
        end
        if !iszero(which)
            push!(momenta, length(unique_momenta))
        end

        i += 1
    end

    return ions, unique_momenta, momenta, indeterminate_momenta, order
end

# 𝐤 is nothing, we want a dipole amplitude
function diagram_amplitude(::Type{Amp}, system::System, diagram::Diagram, iref, i, ::Nothing) where Amp
    amplitude = complex(zero(Amp))

    amplitude
end

# 𝐤 is determinate, we want a photoelectron spectrum
function diagram_amplitude(::Type{Amp}, system::System, diagram::Diagram, iref, i, 𝐤) where Amp
    amplitude = complex(zero(Amp))

    amplitude
end

momentum_type(_, 𝐤) = typeof(𝐤)
momentum_type(system, ::Nothing) = eltype(system.volkov.∫A)

set_momentum!(𝐩s::AbstractVector{<:Number}, 𝐩ₛₜ::Number, i) =
    setindex!(𝐩s, 𝐩ₛₜ, i)
set_momentum!(𝐩s::AbstractVector{<:SVector{3}}, 𝐩ₛₜ::SVector{3}, i) =
    setindex!(𝐩s, 𝐩ₛₜ, i)
set_momentum!(𝐩s::AbstractVector{<:SVector{3}}, 𝐩ₛₜ::T, i) where {T<:Number} =
    setindex!(𝐩s, SVector{3,T}(zero(T), zero(T), 𝐩ₛₜ), i)

function evaluate_momenta!(𝐩s, prefactors, system, unique_momenta, indeterminate_momenta, i; ϵ=1e-2*√(eps(eltype(system.t))))
    for idm in indeterminate_momenta
        uidm = unique_momenta[idm]
        a,b = i[uidm[1]],i[uidm[2]]
        set_momentum!(𝐩s, stationary_momentum(system.volkov, a, b), idm)
        τ = system.t[a]-system.t[b]
        ζ = (2π/(im*τ + ϵ))^(3/2)
        prefactors[idm] = ζ
    end
end

function ionization(system::System, diagram::Diagram, 𝐩, 𝐀, i)
    α,which = diagram.path[end]
    @assert which == 0
    source_term(system.ionization_channels[α],
                i,
                kinematic_momentum(𝐩, 𝐀[i]))
end

function recombination(system::System, diagram::Diagram, 𝐩, 𝐀, i)
    α,which = first(diagram)
    if which == 0 && length(diagram) > 1
        d = system.ionization_channels[α].st.d
        conj(d(kinematic_momentum(𝐩, 𝐀[i])))
    else
        true
    end
end

function integrate_diagram(::Type{Amp}, system::System, diagram::Diagram, iref, 𝐤=nothing; memory=typemax(Int), imin=1) where Amp
    ions, unique_momenta, momenta, indeterminate_momenta, order = analyze_diagram(system, diagram)

    weight = (-im*system.dt)^order

    # println()
    # @info "Integrating diagram up to" iref system diagram ions unique_momenta momenta indeterminate_momenta order weight 𝐤

    # @show
    Eᵢₒₙₛ = [system.ionization_channels[ion].E for ion in ions]
    𝐩s = complex(zeros(momentum_type(system, 𝐤), length(unique_momenta)))
    prefactors = ones(complex(eltype(system.t)), length(unique_momenta))
    if !isnothing(𝐤)
        𝐩s[1] = 𝐤
    end
    # @show 𝐩s

    𝐀 = system.𝐀
    amplitude = complex(zero(Amp))
    ctT = complex(eltype(system.t))

    is = vcat(iref, zeros(Int, order))

    for i in TelescopeIterator(max(1,imin):iref-1, order, memory)
        is[2:end] .= i
        # is = vcat(iref, i)
        # is = (iref,i...)
        # for i in 1:iref-1, j in 1:i-1
        #     is = (iref,i,j)
        # println(is)

        evaluate_momenta!(𝐩s, prefactors, system, unique_momenta, indeterminate_momenta, is)
        # @show 𝐩s prefactors

        Sᵢₒₙ = zero(ctT)
        Sₑₗ = zero(ctT)
        for j = 1:order
            a,b = is[j], is[j+1]
            τ = system.t[a] - system.t[b]
            Sᵢₒₙ += Eᵢₒₙₛ[j]*τ
            Sₑₗ += volkov_phase(𝐩s[j], system.volkov, a, b)
        end
        S = Sᵢₒₙ + Sₑₗ
        aₚᵣₒₚ = prod(prefactors)*exp(-im*S)

        # @show is prefactors

        ∂a = (ionization(system, diagram, 𝐩s[end], 𝐀, is[end]) *
              aₚᵣₒₚ *
              recombination(system, diagram, 𝐩s[1], 𝐀, iref))

        # Loop over "interior" interactions
        for j = (order>1 && first(diagram)[2]==0 ? 3 : 2):order
            ion,which = diagram.path[j-1]
            α = ions[j-1]
            β = ions[j]
            # @show j, ion, which α,β
            interaction = system.couplings[which][α,β]

            𝐤ᵢ = 𝐩s[momenta[j-1]]
            𝐩ᵢ = 𝐩s[momenta[j]]
            𝐀ᵢ = 𝐀[is[j]]

            ∂a *= interaction(kinematic_momentum(𝐤ᵢ, 𝐀ᵢ), kinematic_momentum(𝐩ᵢ, 𝐀ᵢ), is[j+1])
        end

        # ∂a = prod(prefactors)
        amplitude += ∂a
    end

    weight*amplitude
end

# * High-level interface

function photoelectron_spectrum(k::AbstractArray{T},
                                system::System, diagram::Diagram;
                                iref=length(system.t),
                                verbosity=1, kwargs...) where T
    verbosity > 0 && @info "Photoelectrum spectrum calculation" system diagram length(k)

    cT = complex(eltype(T))
    c = similar(k, cT)
    p = Progress(length(k))
    threaded_range_loop(eachindex(k)) do i
        c[i] = integrate_diagram(cT, system, diagram, iref, k[i]; kwargs...)
        ProgressMeter.next!(p)
    end
    c
end

function photoelectron_spectrum(k, args...; kwargs...)
    system = System(args...; kwargs...)
    photoelectron_spectrum(k, system, Diagram(system); kwargs...)
end

"""
    induced_dipole(system, diagram[; kwargs...])

Compute the induced dipole moment as a function of time of the
[`System`](@ref), for a specific [`Diagram`](@ref).
"""
function induced_dipole(system::System, diagram::Diagram; verbosity = 1, kwargs...)
    verbosity > 0 && @info "Induced dipole calculation" system diagram

    t = system.t

    DT = eltype(system.𝐀)
    𝐝 = zeros(DT, length(t))

    memory = get(kwargs, :memory, typemax(Int))

    @showprogress for (i,t) in enumerate(t)
        𝐝̃ = integrate_diagram(DT, system, diagram, i; imin=i-memory, kwargs...)
        𝐝[i] = 2real(𝐝̃)
    end

    𝐝
end

function induced_dipole(args...; kwargs...)
    system = System(args...; kwargs...)
    diagram = Diagram([(1,0),(1,0)], system)
    induced_dipole(system, diagram; kwargs...)
end

# * Exports

export Diagram, photoelectron_spectrum, induced_dipole


text:
@defstruct ConvolutionLayer Layer (
  (name :: AbstractString = "", !isempty(name)),
  param_key :: AbstractString = "",
  (bottoms :: Vector{Symbol} = Symbol[], length(bottoms) > 0),
  (tops :: Vector{Symbol} = Symbol[], length(tops) == length(bottoms)),
  (kernel :: NTuple{2,Int} = (1,1), all([kernel...] .> 0)),
  (stride :: NTuple{2,Int} = (1,1), all([stride...] .> 0)),
  (pad :: NTuple{2, Int} = (0,0), all([pad...] .>= 0)),
  (n_filter :: Int = 1, n_filter > 0),
  (n_group :: Int = 1, n_group > 0),
  neuron :: ActivationFunction = Neurons.Identity(),
  filter_init :: Initializer = XavierInitializer(),
  bias_init :: Initializer = ConstantInitializer(0),
  filter_regu :: Regularizer = L2Regu(1),
  bias_regu :: Regularizer = NoRegu(),
  filter_cons :: Constraint = NoCons(),
  bias_cons :: Constraint = NoCons(),
  filter_lr :: AbstractFloat = 1.0,
  bias_lr :: AbstractFloat = 2.0,
)
@characterize_layer(ConvolutionLayer,
  has_param  => true,
  has_neuron => true,
  can_do_bp  => true
)

type CPUConvState
  col_buffer      :: Blob
  M               :: Int
  N               :: Int
  K               :: Int

  # a all-1 vector used in gemm to help bias calculation
  bias_multiplier :: Blob
  img_buffer      :: Array
end

function setup_etc(backend::CPUBackend, layer::ConvolutionLayer, dtype, width, height,
    channels, batch_size, width_out, height_out, inputs)

  if layer.kernel[1] == 1 && layer.kernel[2] == 1 &&
     layer.stride[1] == 1 && layer.stride[2] == 1 &&
     layer.pad[1] == 0 && layer.pad[2] == 0
    col_buffer = NullBlob()
  else
    col_buffer = CPUBlob(Array(dtype, width_out, height_out, channels*prod(layer.kernel), 1))
  end
  M = height_out * width_out
  N = div(layer.n_filter, layer.n_group)
  K = div(channels * layer.kernel[1] * layer.kernel[2], layer.n_group)
  bias_multiplier = make_blob(backend, dtype, M, 1, 1, 1)
  fill!(bias_multiplier, convert(dtype,1))
  img_buffer = Array(dtype, width, height, channels)
  etc = CPUConvState(col_buffer, M, N, K, bias_multiplier, img_buffer)
  return etc
end

type ConvolutionLayerState <: LayerState
  layer      :: ConvolutionLayer
  blobs      :: Vector{Blob}
  blobs_diff :: Vector{Blob}

  parameters :: Vector{Parameter}

  filter  :: Blob
  ∇filter :: Blob
  bias    :: Blob
  ∇bias   :: Blob

  ConvolutionLayerState(backend::Backend, layer::ConvolutionLayer, shared_params, inputs::Vector{Blob}) = begin
    for i = 1:length(inputs)
      # currently we only handle 4D-tensor
      @assert ndims(inputs[i]) == 4
    end

    width, height, channels, batch_size = size(inputs[1])
    @assert channels % layer.n_group == 0
    @assert layer.n_filter % layer.n_group == 0

    width_out  = div(width  + 2*layer.pad[1]-layer.kernel[1], layer.stride[1]) + 1
    height_out = div(height + 2*layer.pad[2]-layer.kernel[2], layer.stride[2]) + 1

    dtype = eltype(inputs[1])

    # Make sure all input blobs are of the same shape
    for i = 2:length(inputs)
      @assert (width,height,channels,batch_size) == size(inputs[i])
      @assert dtype == eltype(inputs[i])
    end

    blobs = Array(Blob, length(inputs))
    blobs_diff = Array(Blob, length(inputs))

    for i = 1:length(inputs)
      blobs[i] = make_blob(backend, dtype, width_out, height_out, layer.n_filter, batch_size)
      blobs_diff[i] = make_blob(backend, dtype, width_out, height_out, layer.n_filter, batch_size)
    end

    if shared_params != nothing
      @assert length(shared_params) == 2
      @assert shared_params[1].name == "filter" && shared_params[2].name == "bias"
      @assert size(shared_params[1].blob) == tuple(layer.kernel...,div(channels,layer.n_group),layer.n_filter)
      @assert eltype(shared_params[1].blob) == dtype
      @debug("ConvolutionLayer($(layer.name)): sharing filters and bias")

      param_filter, param_bias = [share_parameter(backend, param) for param in shared_params]
    else
      param_filter = make_parameter(backend,"filter",dtype,(layer.kernel[1],layer.kernel[2],div(channels,layer.n_group), layer.n_filter),
          layer.filter_init, layer.filter_regu, layer.filter_cons, layer.filter_lr)
      param_bias   = make_parameter(backend,"bias",dtype,(layer.n_filter,),
          layer.bias_init, layer.bias_regu, layer.bias_cons, layer.bias_lr)
    end

    filter     = param_filter.blob
    ∇filter    = param_filter.gradient
    bias       = param_bias.blob
    ∇bias      = param_bias.gradient
    parameters = [param_filter, param_bias]

    etc = setup_etc(backend, layer, dtype, width, height, channels, batch_size, width_out, height_out, inputs)

    state = new(layer, blobs, blobs_diff, parameters)
    state.filter = filter
    state.∇filter = ∇filter
    state.bias = bias
    state.∇bias = ∇bias

    state.height_out = height_out
    state.width_out = width_out

    state.etc = etc

    state.frozen = false

    return state
  end

  # Auxiliary variables
  height_out :: Int
  width_out  :: Int

  etc        :: Any # whatever status a computation backend needs to maintain

  frozen     :: Bool
end

function freeze!(state::ConvolutionLayerState)
  state.frozen = true
end
function unfreeze!(state::ConvolutionLayerState)
  state.frozen = false
end
function is_frozen(state::ConvolutionLayerState)
  state.frozen
end

function setup(backend::Backend, layer::ConvolutionLayer, shared_state, inputs::Vector{Blob}, diffs::Vector{Blob})
  return ConvolutionLayerState(backend, layer, shared_state, inputs)
end
function setup(backend::Backend, layer::ConvolutionLayer, inputs::Vector{Blob}, diffs::Vector{Blob})
  setup(backend, layer, nothing, inputs, diffs)
end
function shutdown_etc(backend::CPUBackend, state::ConvolutionLayerState)
end
function shutdown(backend::Backend, state::ConvolutionLayerState)
  map(destroy, state.blobs)
  map(destroy, state.blobs_diff)
  map(destroy, state.parameters)

  shutdown_etc(backend, state)
end

function forward(backend::CPUBackend, state::ConvolutionLayerState, inputs::Vector{Blob})
  for i = 1:length(inputs)
    input = inputs[i]
    output = state.blobs[i]
    erase!(output)
    dtype = eltype(input)
    width, height, channels, num = size(input)
    img_offset = width*height*channels * sizeof(dtype)
    weight_offset = state.etc.N * state.etc.K * sizeof(dtype)
    col_offset = state.etc.M * state.etc.K * sizeof(dtype)
    top_offset = state.etc.M * state.etc.N * sizeof(dtype)
    top_img_offset = state.height_out * state.width_out * state.layer.n_filter * sizeof(dtype)

    for n = 1:num
      if isa(state.etc.col_buffer, NullBlob)
        col_buffer = pointer(input.data) + img_offset * (n-1)
      else
        col_buffer = state.etc.col_buffer.data
        im2col(input.data, n, col_buffer,
            width, height, channels, state.layer.kernel, state.layer.pad, state.layer.stride)
        col_buffer = pointer(col_buffer)
      end

      output_ptr = pointer(output.data) + top_img_offset * (n-1)
      for g = 1:state.layer.n_group
        RawBLAS.gemm!('N', 'N', state.etc.M, state.etc.N, state.etc.K, convert(dtype, 1),
            col_buffer + col_offset * (g-1),
            convert(Ptr{dtype}, pointer(state.filter.data)) + weight_offset * (g-1),
            convert(dtype, 0), output_ptr + top_offset * (g-1))
      end
      RawBLAS.gemm!('N', 'N', state.etc.M, state.layer.n_filter, 1, convert(dtype, 1),
          state.etc.bias_multiplier.data, pointer(state.bias.data), convert(dtype, 1), output_ptr)
    end
  end
end

function backward(backend::CPUBackend, state::ConvolutionLayerState, inputs::Vector{Blob}, diffs::Vector{Blob})
  erase!(state.∇filter)
  erase!(state.∇bias)
  for i = 1:length(inputs)
    input = inputs[i]
    top_diff = state.blobs_diff[i]
    dtype = eltype(input)
    width, height, channels, num = size(input)
    img_offset = width*height*channels * sizeof(dtype)
    weight_offset = state.etc.N * state.etc.K * sizeof(dtype)
    col_offset = state.etc.M * state.etc.K * sizeof(dtype)
    top_offset = state.etc.M * state.etc.N * sizeof(dtype)
    top_img_offset = state.height_out * state.width_out * state.layer.n_filter * sizeof(dtype)

    if !state.frozen
      for n = 1:num
        top_diff_ptr = pointer(top_diff.data) + top_img_offset * (n-1)

        #----------------------------------------------
        # bias gradient
        RawBLAS.gemv!('T', state.etc.M, state.layer.n_filter, convert(dtype, 1), top_diff_ptr,
            state.etc.bias_multiplier.data, convert(dtype, 1), pointer(state.∇bias.data))

        #----------------------------------------------
        # filter gradient
        if isa(state.etc.col_buffer, NullBlob)
          col_buffer = pointer(input.data) + img_offset * (n-1)
        else
          col_buffer = state.etc.col_buffer.data
          im2col(input.data, n, col_buffer,
              width, height, channels, state.layer.kernel, state.layer.pad, state.layer.stride)
          col_buffer = pointer(col_buffer)
        end
        for g = 1:state.layer.n_group
          RawBLAS.gemm!('T', 'N', state.etc.K, state.etc.N, state.etc.M, convert(dtype, 1),
              col_buffer + col_offset * (g-1),
              top_diff_ptr + top_offset * (g-1), convert(dtype, 1),
              convert(Ptr{dtype}, pointer(state.∇filter.data)) + weight_offset * (g-1))
        end
      end
    end

    #----------------------------------------------
    # back propagate gradient
    if isa(diffs[i], CPUBlob)
      diff = diffs[i]
      for n = 1:num
        top_diff_ptr = pointer(top_diff.data) + top_img_offset * (n-1)
        if isa(state.etc.col_buffer, NullBlob)
          col_buffer = pointer(diff.data) + img_offset * (n-1)
        else
          col_buffer = pointer(state.etc.col_buffer.data)
        end

        for g = 1:state.layer.n_group
          RawBLAS.gemm!('N', 'T', state.etc.M, state.etc.K, state.etc.N, convert(dtype, 1),
              top_diff_ptr + top_offset * (g-1),
              pointer(state.filter.data) + weight_offset * (g-1),
              convert(dtype, 0), col_buffer + col_offset * (g-1))
        end
        if !(isa(state.etc.col_buffer, NullBlob))
           col2im(state.etc.col_buffer.data, diff.data, n, state.etc.img_buffer,
              width, height, channels, state.layer.kernel, state.layer.pad, state.layer.stride)
        end
      end
    end
  end
end


text:
type ALSWR{TP<:Parallelism,TI<:Inputs,TM<:Model}
    inp::TI
    model::Nullable{TM}
    par::TP
end

ALSWR(inp::FileSpec, par::ParShmem) = ALSWR{ParShmem,SharedMemoryInputs,SharedMemoryModel}(SharedMemoryInputs(inp), nothing, par)
ALSWR(user_item_ratings::FileSpec, item_user_ratings::FileSpec, par::ParBlob) = ALSWR{ParBlob,DistInputs,DistModel}(DistInputs(user_item_ratings, item_user_ratings), nothing, par)

function clear(als::ALSWR)
    clear(als.inp)
    isnull(als.model) || clear(get(als.model))
end

function localize!(als::ALSWR)
    localize!(als.inp)
    isnull(als.model) || localize!(get(als.model))
end

function save(als::ALSWR, filename::AbstractString)
    clear(als)
    open(filename, "w") do f
        serialize(f, als)
    end
    nothing
end

function train(als::ALSWR, niters::Int, nfacts::Int64, lambda::Float64=0.065)
    als.model = prep(als.inp, nfacts, lambda)
    fact_iters(als, niters)
    nothing
end

fact_iters(als, niters) = fact_iters(als.par, get(als.model), als.inp, niters)

##
# Training
#
function update_user(u::Int64)
    c = fetch_compdata()
    update_user(u, c.model, c.inp, c.lambdaI)
end

function update_user(r::UnitRange)
    c = fetch_compdata()
    for u in r
        update_user(u, c.model, c.inp, c.lambdaI)
    end
end

function update_user(u::Int64, model, inp, lambdaI)
    nzrows, nzvals = items_and_ratings(inp, u)
    Pu = getP(model, nzrows)
    vec = Pu * nzvals
    mat = (Pu * Pu') + (length(nzrows) * lambdaI)
    setU(model, u, mat \ vec)
    nothing
end

function update_item(r::UnitRange)
    c = fetch_compdata()
    for i in r
        update_item(i::Int64, c.model, c.inp, c.lambdaI)
    end
    nothing
end

function update_item(i::Int64)
    c = fetch_compdata()
    update_item(i::Int64, c.model, c.inp, c.lambdaI)
end

function update_item(i::Int64, model, inp, lambdaI)
    nzrows, nzvals = users_and_ratings(inp, i)
    Uit = getU(model, nzrows)
    Ui = Uit'
    vec = Uit * nzvals
    mat = (Uit * Ui) + (length(nzrows) * lambdaI)
    setP(model, i, mat \ vec)
    nothing
end

##
# Validation
function rmse_block_sz(als::ALSWR)
    model = get(als.model)
    # heuristics to get a reasonable split size (applicable only to single node system)
    block_sz = floor(Int, Base.Sys.free_memory() / 2 / nworkers() / nitems(model) / sizeof(Float64))
    block_sz = min(5000, block_sz)
    @logmsg("chosen block size: $block_sz. freemem: $(Base.Sys.free_memory()), nworkers: $(nworkers()), nitems: $(nitems(model))")
    block_sz
end

function rmse(als::ALSWR; block_sz::Int=rmse_block_sz(als))
    ensure_loaded(als.inp)
    rmse(als, als.inp; block_sz=block_sz)
end

function rmse(als::ALSWR, testdataset::FileSpec; block_sz::Int=rmse_block_sz(als))
    ensure_loaded(als.inp)
    i_idmap = item_idmap(als.inp)

    testinp = SharedMemoryInputs(testdataset)
    ensure_loaded(testinp; only_items=i_idmap)
    rmse(als, testinp; block_sz=block_sz)
end

##
# Recommendation
function _recommend(Uvec, model, rated, i_idmap; count::Int=10)
    top = sortperm(vec(mul_p(model, Uvec)))

    recommended = Int64[]
    idx = 1
    while length(recommended) < count && length(top) >= idx
        item_id = top[idx]
        real_item_id = isempty(i_idmap) ? item_id : i_idmap[item_id]
        (item_id in rated) || push!(recommended, real_item_id)
        idx += 1
    end
    nexcl = idx - count - 1

    mapped_rated = isempty(i_idmap) ? rated : Int64[i_idmap[id] for id in rated]
    recommended, mapped_rated, nexcl
end

function recommend(als::ALSWR, user::Int; unrated::Bool=true, count::Int=10)
    ensure_loaded(als.inp)
    i_idmap = item_idmap(als.inp)
    u_idmap = user_idmap(als.inp)

    if !isempty(u_idmap)
        (user in u_idmap) || (return (Int[], Int[], 0))
        user = findfirst(u_idmap, user)
    end

    model = get(als.model)

    # All the items sorted in decreasing order of rating.
    Uvec = reshape(getU(model, user), 1, nfactors(model))
    rated = unrated ? all_items_rated(als.inp, user) : Int64[]

    _recommend(Uvec, model, rated, i_idmap; count=count)
end

function recommend(als::ALSWR, user_ratings::SparseVector{Float64,Int64}; unrated::Bool=true, count::Int=10)
    # R = U * P
    # given a new row in R, figure out the corresponding row in U
    #
    # Rvec = Uvec * P
    # Uvec = Rvec * Pinv
    #
    # since: I = (P * Pt) * inv(P * Pt)
    # Pinv = Pt * inv(P * Pt)
    #
    # Uvec = Rvec * (Pt * inv(P * Pt))
    ensure_loaded(als.inp)
    i_idmap = item_idmap(als.inp)
    model = get(als.model)
    mapped_ratings = isempty(i_idmap) ? full(user_ratings) : user_ratings[i_idmap]
    Rvec = reshape(mapped_ratings, 1, length(mapped_ratings))
    Uvec = mul_pinv(model, Rvec)

    _recommend(Uvec, model, find(Rvec), i_idmap; count=count)
end

##
# Shared memory parallelism
type ComputeData{TM<:Model,TI<:Inputs}
    model::TM
    inp::TI
    lambdaI::ModelFactor # store directly, avoid null check on every iteration
end

const compdata = ComputeData[]

function share_compdata(c::ComputeData)
    push!(compdata, c)
    ensure_loaded(c.inp)
    ensure_loaded(c.model)
    nothing
end
fetch_compdata() = compdata[1]
noop(args...) = nothing

function fact_iters{TP<:ParShmem,TM<:Model,TI<:Inputs}(::TP, model::TM, inp::TI, niters::Int64)
    t1 = time()
    share!(model)
    share!(inp)

    c = ComputeData(model, inp, get(model.lambdaI))
    @sync for w in workers()
        @async remotecall_fetch(share_compdata, w, c)
    end

    nu = nusers(inp)
    ni = nitems(inp)
    @logmsg("nusers: $nu, nitems: $ni")
    for iter in 1:niters
        @logmsg("begin iteration $iter")
        #pmap(update_user, 1:nu)
        @parallel (noop) for u in 1:nu
            update_user(u)
        end
        @logmsg("\tusers")
        #pmap(update_item, 1:ni)
        @parallel (noop) for i in 1:ni
            update_item(i)
        end
        @logmsg("\titems")
    end

    localize!(model)

    t2 = time()
    @logmsg("fact time $(t2-t1)")
    nothing
end

function rmse(inp::Inputs, model::Model, r::UnitRange{Int})
    Upart = getU(model, r) # nfactors x block_sz
    # Upart' is block_sz x nfactors
    # P is nfactors x nitems
    # predicted is of size block_sz x nitems
    predicted = mul_p(model, Upart')

    cumerr = 0
    nvals = 0
    uoffset = first(r)-1

    for user in r
        nzrows, nzvals = items_and_ratings(inp, user)
        nvals += length(nzrows)
        cumerr += sum((vec(predicted[user-uoffset,nzrows]) .- nzvals) .^ 2)
    end
    [cumerr, nvals]
end

function rmse{T<:Union{ParShmem,ParBlob},TI<:Inputs}(als::ALSWR{T}, inp::TI; block_sz::Int=rmse_block_sz(als))
    t1 = time()

    model = get(als.model)
    share!(model)
    share!(inp)
    ensure_loaded(inp)
    ensure_loaded(model)

    nU = nusers(model)
    usplits = [1:block_sz:nU...]
    usplits[end] = nU
    uranges = UnitRange{Int}[usplits[x-1]:usplits[x] for x in 2:length(usplits)]
    cumerr = @parallel (.+) for r in uranges
        rmse(inp, model, r)
    end
    localize!(model)
    @logmsg("rmse time $(time()-t1)")
    sqrt(cumerr[1]/cumerr[2])
end

##
# Blob based distributed memory parallelism
function train(als::ALSWR{ParBlob,DistInputs,DistModel}, niters::Int, nfacts::Int64, model_dir::AbstractString, maxcache::Int=def_cache(), lambda::Float64=0.065)
    als.model = prep(als.inp, nfacts, lambda, model_dir, maxcache)
    fact_iters(als, niters)
    nothing
end

function update_user_blobs(r::UnitRange)
    c = fetch_compdata()
    U = get(c.model.U)
    flush(U; callback=false)
    update_user(r)
    save(U)
    nothing
end

function update_item_blobs(r::UnitRange)
    c = fetch_compdata()
    P = get(c.model.P)
    flush(P; callback=false)
    update_item(r)
    save(P)
    nothing
end

function fact_iters{TP<:ParBlob,TM<:Model,TI<:Inputs}(::TP, model::TM, inp::TI, niters::Int64)
    t1 = time()

    clear(inp)
    share!(model)
    share!(inp)

    c = ComputeData(model, inp, get(model.lambdaI))
    uranges = UnitRange[p.first for p in get(model.U).splits]
    iranges = UnitRange[p.first for p in get(model.P).splits]

    # clear, share the data and load it again (not required, but more efficient)
    clear(model)
    W = workers()
    @sync for w in W
        @async remotecall_fetch(share_compdata, w, c)
    end
    ensure_loaded(model)
    U = get(model.U)
    P = get(model.P)

    nu = nusers(inp)
    ni = nitems(inp)
    @logmsg("nusers: $nu, nitems: $ni")
    for iter in 1:niters
        @logmsg("begin iteration $iter")
        pmap(update_user_blobs, uranges)
        @logmsg("\tusers")
        pmap(update_item_blobs, iranges)
        @logmsg("\titems")
    end

    localize!(model)

    t2 = time()
    @logmsg("fact time $(t2-t1)")
    nothing
end

##
# Thread parallelism
if (Base.VERSION >= v"0.5.0-")

ALSWR(inp::FileSpec, par::ParThread) = ALSWR{ParThread,SharedMemoryInputs,SharedMemoryModel}(SharedMemoryInputs(inp), nothing, par)

function thread_update_item{TM<:Model,TI<:Inputs}(model::TM, inp::TI, ni::Int64, lambdaI::Matrix{Float64})
    @threads for i in Int64(1):ni
        update_item(i, model, inp, lambdaI)
    end
    nothing
end

function thread_update_user{TM<:Model,TI<:Inputs}(model::TM, inp::TI, nu::Int64, lambdaI::Matrix{Float64})
    @threads for u in Int64(1):nu
        update_user(u, model, inp, lambdaI)
    end
    nothing
end
function fact_iters{TP<:ParThread,TM<:Model,TI<:Inputs}(::TP, model::TM, inp::TI, niters::Int64)
    t1 = time()

    nu = nusers(inp)
    ni = nitems(inp)
    lambdaI = get(model.lambdaI)
    for iter in 1:niters
        @logmsg("begin iteration $iter")
        # gc is not threadsafe yet. issue #10317
        thread_update_user(model, inp, nu, lambdaI)
        @logmsg("\tusers")
        thread_update_item(model, inp, ni, lambdaI)
        @logmsg("\titems")
    end

    t2 = time()
    @logmsg("fact time $(t2-t1)")
    nothing
end

function rmse{TP<:ParThread,TI<:Inputs}(als::ALSWR{TP}, inp::TI; block_sz::Int=0)
    t1 = time()

    model = get(als.model)
    errs = zeros(nthreads())
    lengths = zeros(Int, nthreads())

    pos = 1
    NF = nfactors(model)
    N2 = nusers(als.inp)
    while pos < N2
        endpos = min(pos+10000, N2)
        @threads for user in pos:endpos
            Uvec = reshape(getU(model, user), 1, NF)
            nzrows, nzvals = items_and_ratings(inp, user)
            predicted = vec(mul_p(model, Uvec))[nzrows]

            tid = threadid()
            lengths[tid] += length(predicted)
            errs[tid] += sum((predicted - nzvals) .^ 2)
        end
        pos = endpos + 1
    end
    @logmsg("rmse time $(time()-t1)")
    sqrt(sum(errs)/sum(lengths))
end
end


text:
include("wannierize.jl")

filename = "silicon"

read_amn = false
read_eig = false

p = read_system(filename,read_amn,read_eig)

nbeg = 1
nend = p.nwannier

#Set the seed of the pseudo-random generator
#srand(0)

println("$(p.nband) bands, wannierizing bands $nbeg:$nend, $(p.N1) x $(p.N2) x $(p.N3) grid, $(p.nntot) neighbors")

#Build the gauge that makes the Bloch frame continuous on the Brillouin Zone.
#This is equivalent to building a set of algebraic decaying Wannier functions
p, Obs_array, Uint = make_wannier(p)

# #Plot and display the results
# plot_results(p)
# print_error(p)

# figure()
# for i=1:5:p.N2
#     plot(real(Uint[:,i,2,2]),imag(Uint[:,i,2,2]),"x-")
# end
# title("real 22 / imag 22")

# figure()
# for i=1:10:p.N2
#     plot(real(Uint[:,i,1,2]),imag(Uint[:,i,1,2]),"x-")
# end
# title("real 12 / imag 12")

# figure()
# for i=1:5:p.N2
#     plot(real(Uint[:,i,:,2]))
# end

# #using Plots
# #println("Creating animation of the obstruction")
# #anim = @animate for i=1:2*p.N2-1
# #               plot(real(Uint[:,abs(p.N2-i)+1,2,2]),imag(Uint[:,abs(p.N2-i)+1,2,2]),xlims=[-1,1],ylims=[-1,1])
# #end
# #
# #gif(anim,"obs.gif",fps=25)

max_stars_repo_name: VVitale/wannier

max_issues_repo_name: VVitale/wannier

max_forks_repo_name: VVitale/wannier


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

module Dir

import Pkg
import ..DEFAULT_META, ..META_BRANCH, ..PkgError
import LibGit2, LibGit2.with
const DIR_NAME = ".julia"

_pkgroot() = abspath(get(ENV,"JULIA_PKGDIR",joinpath(homedir(),DIR_NAME)))
isversioned(p::AbstractString) = ((x,y) = (VERSION.major, VERSION.minor); basename(p) == "v$x.$y")

function path()
    b = _pkgroot()
    x, y = VERSION.major, VERSION.minor
    d = joinpath(b,"v$x.$y")
    if isdir(d) || !isdir(b) || !isdir(joinpath(b, "METADATA"))
        return d
    end
    return b
end
path(pkg::AbstractString...) = normpath(path(),pkg...)

function cd(f::Function, args...; kws...)
    dir = path()
    metadata_dir = joinpath(dir, "METADATA")
    if !isdir(metadata_dir)
        !haskey(ENV,"JULIA_PKGDIR") ? init() :
            throw(PkgError("Package metadata directory $metadata_dir doesn't exist; run Pkg.init() to initialize it."))
    end
    if haskey(ENV,"JULIA_PKGDIR")
        withenv("JULIA_PKGDIR" => abspath(ENV["JULIA_PKGDIR"])) do
            Base.cd(()->f(args...; kws...), dir)
        end
    else
        Base.cd(()->f(args...; kws...), dir)
    end
end

function init(meta::AbstractString=DEFAULT_META, branch::AbstractString=META_BRANCH)
    dir = path()
    @info "Initializing package repository $dir"
    metadata_dir = joinpath(dir, "METADATA")
    if isdir(metadata_dir)
        @info "Package directory $dir is already initialized"
        LibGit2.set_remote_url(metadata_dir, "origin", meta)
        return
    end
    local temp_dir = ""
    try
        mkpath(dir)
        temp_dir = mktempdir(dir)
        Base.cd(temp_dir) do
            @info "Cloning METADATA from $meta"
            with(LibGit2.clone(meta, "METADATA", branch = branch)) do metadata_repo
                LibGit2.set_remote_url(metadata_repo, "origin", meta)
            end
            touch("REQUIRE")
            touch("META_BRANCH")
            write("META_BRANCH", branch)
        end
        #Move TEMP to METADATA
        Base.mv(joinpath(temp_dir,"METADATA"), metadata_dir)
        Base.mv(joinpath(temp_dir,"REQUIRE"), joinpath(dir,"REQUIRE"))
        Base.mv(joinpath(temp_dir,"META_BRANCH"), joinpath(dir,"META_BRANCH"))
        rm(temp_dir, recursive=true)
    catch err
        ispath(metadata_dir) && rm(metadata_dir, recursive=true)
        ispath(temp_dir) && rm(temp_dir, recursive=true)
        rethrow(err)
    end
end

function getmetabranch()
    try
        readline(joinpath(path(),"META_BRANCH"))
    catch err
        META_BRANCH
    end
end

end # module

max_stars_repo_name: djsegal/julia-fork

max_issues_repo_name: djsegal/julia-fork

max_forks_repo_name: djsegal/julia-fork


text:
using Weber
using WeberCedrus
using Base.Test

# these tests are copied from Weber, and merely verify that ane expeirment still
# runs properly with the extension installed.

# this allows a test of timing for an experiment without setting up any
# multimedia resources, so it can be run just about anywhere.
function find_timing(fn;keys...)
  empty!(Weber.null_record)
  exp = Experiment(;null_window=true,hide_output=true,extensions=[CedrusXID()],keys...)
  setup(() -> fn(),exp)
  run(exp,await_input=false)

  nostarts = filter(x -> !endswith(string(x[:code]),"_start") &&
                         x[:code] != "terminated",
                    Weber.null_record)
  map(x -> x[:code],nostarts),map(x -> x[:time],nostarts),nostarts
end

seq_events,_,_ = find_timing() do
  addtrial(moment(0.01,() -> record(:a)),
           moment(0.01,() -> record(:b)),
           moment(0.01,() -> record(:c)))
end

@testset "Moment Sequencing" begin
  @test seq_events == [:a,:b,:c]
end


text:
using ProgressMeter
using Random

function prepare_runner()
    cd(joinpath(dirname(@__DIR__), "runner")) do
        Base.run(`docker build . -t newpkgeval`)
    end
    return
end

"""
    run_sandboxed_julia(julia::VersionNumber, args=``; wait=true, interactive=true,
                        stdin=stdin, stdout=stdout, stderr=stderr, kwargs...)

Run Julia inside of a sandbox, passing the given arguments `args` to it. The argument
`julia` specifies the version of Julia to use.

The argument `wait` determines if the process will be waited on, and defaults to true. If
setting this argument to `false`, remember that the sandbox is using on Docker and killing
the process does not necessarily kill the container. It is advised to use the `name` keyword
argument to set the container name, and use that to kill the Julia process.

The keyword argument `interactive` maps to the Docker option, and defaults to true.
"""
function run_sandboxed_julia(julia::VersionNumber, args=``; wait=true,
                             stdin=stdin, stdout=stdout, stderr=stderr, kwargs...)
    container = spawn_sandboxed_julia(julia, args; kwargs...)
    Base.run(pipeline(`docker attach $container`, stdin=stdin, stdout=stdout, stderr=stderr); wait=wait)
end

function spawn_sandboxed_julia(julia::VersionNumber, args=``; interactive=true, name=nothing)
    cmd = `docker run --detach`

    # mount data
    @assert ispath(julia_path(julia))
    installed_julia_path = installed_julia_dir(julia)
    @assert isdir(installed_julia_path)
    registry_path = joinpath(first(DEPOT_PATH), "registries")
    @assert isdir(registry_path)
    cmd = ```$cmd --mount type=bind,source=$installed_julia_path,target=/opt/julia,readonly
                  --mount type=bind,source=$registry_path,target=/root/.julia/registries,readonly```

    if interactive
        cmd = `$cmd --interactive --tty`
    end

    if name !== nothing
        cmd = `$cmd --name $name`
    end

    container = chomp(read(`$cmd --rm newpkgeval /opt/julia/bin/julia $args`, String))
    return something(name, container)
end

"""
    run_sandboxed_test(julia::VersionNumber, pkg; do_depwarns=false, log_limit=5*1024^2,
                       time_limit=60*45)

Run the unit tests for a single package `pkg` inside of a sandbox using Julia version
`julia`. If `do_depwarns` is `true`, deprecation warnings emitted while running the
package's tests will cause the tests to fail. Test will be forcibly interrupted after
`time_limit` seconds or if the log becomes larger than `log_limit`.

A log for the tests is written to a version-specific directory in the NewPkgEval root
directory.

Refer to `run_sandboxed_julia`[@ref] for more possible `keyword arguments.
"""
function run_sandboxed_test(julia::VersionNumber, pkg; log_limit = 5*1024^2 #= 5 MB =#,
                            time_limit = 45*60, do_depwarns=false, kwargs...)
    # everything related to testing in Julia: version compatibility, invoking Pkg, etc

    if pkg.name in skip_lists[pkg.registry]
        return missing, :skip, :explicit, missing
    elseif endswith(pkg.name, "_jll")
        return missing, :skip, :jll, missing
    end

    # can we even test this package?
    pkg_versions = Pkg.Operations.load_versions(pkg.path)
    julia_supported = Dict{VersionNumber,Bool}()
    pkg_compat = Pkg.Operations.load_package_data_raw(Pkg.Types.VersionSpec,
                                                      joinpath(pkg.path, "Compat.toml"))
    for (version_range, bounds) in pkg_compat
        if haskey(bounds, "julia")
            for pkg_version in keys(pkg_versions)
                if pkg_version in version_range
                    julia_supported[pkg_version] = julia ∈ bounds["julia"]
                end
            end
        end
    end
    if length(julia_supported) != length(pkg_versions)
        # not all versions have a bound for Julia, so we need to be conservative
        supported = true
    else
        supported = any(values(julia_supported))
    end
    if !supported
        return missing, :skip, :unsupported, missing
    end

    # prepare for launching a container
    container = "Julia_v$(julia)-$(pkg.name)"
    arg = """
        using Pkg

        # Prevent Pkg from updating registy on the Pkg.add
        Pkg.UPDATED_REGISTRY_THIS_SESSION[] = true

        ENV["CI"] = true
        ENV["PKGEVAL"] = true

        Pkg.add($(repr(pkg.name)))
        Pkg.test($(repr(pkg.name)))
    """
    cmd = do_depwarns ? `--depwarn=error` : ``
    cmd = `$cmd -e $arg`

    mktemp() do path, f
        p = run_sandboxed_julia(julia, cmd; stdout=f, stderr=f, stdin=devnull,
                                interactive=false, wait=false, name=container, kwargs...)
        status = nothing
        reason = missing
        version = missing

        # kill on timeout
        t = Timer(time_limit) do timer
            process_running(p) || return
            status = :kill
            reason = :time_limit
            kill_container(p, container)
        end

        # kill on too-large logs
        t2 = @async while true
            process_running(p) || break
            if stat(path).size > log_limit
                kill_container(p, container)
                status = :kill
                reason = :log_limit
                break
            end
            sleep(2)
        end

        succeeded = success(p)
        log = read(path, String)
        close(t)
        wait(t2)

        # pick up the installed package version from the log
        let match = match(Regex("Installed $(pkg.name) .+ v(.+)"), log)
            if match !== nothing
                version = VersionNumber(match.captures[1])
            end
        end

        if succeeded
            @assert status == nothing
            status = :ok
        elseif status === nothing
            status = :fail
            reason = :unknown

            # figure out a more accurate failure reason from the log
            if occursin("ERROR: Unsatisfiable requirements detected for package", log)
                # NOTE: might be the package itself, or one of its dependencies
                reason = :unsatisfiable
            elseif occursin("ERROR: Package $(pkg.name) did not provide a `test/runtests.jl` file", log)
                reason = :untestable
            elseif occursin("cannot open shared object file: No such file or directory", log)
                reason = :binary_dependency
            elseif occursin(r"Package .+ does not have .+ in its dependencies", log)
                reason = :missing_dependency
            elseif occursin("Some tests did not pass", log)
                reason = :test_failures
            elseif occursin("ERROR: LoadError: syntax", log)
                reason = :syntax
            end
        end

        return version, status, reason, log
    end
end

function kill_container(p, container)
    Base.run(`docker kill --signal=SIGTERM $container`)
    sleep(3)
    if process_running(p)
        Base.run(`docker kill --signal=SIGKILL $container`)
    end
end

function run(julia_versions::Vector{VersionNumber}, pkgs::Vector;
             ninstances::Integer=Sys.CPU_THREADS, callback=nothing, kwargs...)
    # here we deal with managing execution: spawning workers, output, result I/O, etc

    jobs = [(julia=julia, pkg=pkg) for julia in julia_versions for pkg in pkgs]

    # use a random test order to (hopefully) get a more reasonable ETA
    shuffle!(jobs)

    njobs = length(jobs)
    ninstances = min(njobs, ninstances)
    running = Vector{Union{Nothing, eltype(jobs)}}(nothing, ninstances)
    times = DateTime[now() for i = 1:ninstances]
    all_workers = Task[]

    done = false
    did_signal_workers = false
    function stop_work()
        if !done
            done = true
            if !did_signal_workers
                for (i, task) in enumerate(all_workers)
                    task == current_task() && continue
                    Base.istaskdone(task) && continue
                    try; schedule(task, InterruptException(); error=true); catch; end
                    running[i] = nothing
                end
                did_signal_workers = true
            end
        end
    end

    start = now()
    io = IOContext(IOBuffer(), :color=>true)
    on_ci = parse(Bool, get(ENV, "CI", "false"))
    p = Progress(njobs; barlen=50, color=:normal)
    function update_output()
        # known statuses
        o = count(==(:ok),      values(result))
        s = count(==(:skip),    values(result))
        f = count(==(:fail),    values(result))
        k = count(==(:kill),    values(result))
        # remaining
        x = njobs - length(result)

        function runtimestr(start)
            time = Dates.canonicalize(Dates.CompoundPeriod(now() - start))
            if isempty(time.periods) || first(time.periods) isa Millisecond
                "just started"
            else
                # be coarse in the name of brevity
                string(first(time.periods))
            end
        end

        if on_ci
            println("$x packages to test ($o succeeded, $f failed, $k killed, $s skipped, $(runtimestr(start)))")
            sleep(10)
        else
            print(io, "Success: ")
            printstyled(io, o; color = :green)
            print(io, "\tFailed: ")
            printstyled(io, f; color = Base.error_color())
            print(io, "\tKilled: ")
            printstyled(io, k; color = Base.error_color())
            print(io, "\tSkipped: ")
            printstyled(io, s; color = Base.warn_color())
            println(io, "\tRemaining: ", x)
            for i = 1:ninstances
                job = running[i]
                str = if job === nothing
                    " #$i: -------"
                else
                    " #$i: $(job.pkg.name) @ $(job.julia) ($(runtimestr(times[i])))"
                end
                if i%2 == 1 && i < ninstances
                    print(io, rpad(str, 50))
                else
                    println(io, str)
                end
            end
            print(String(take!(io.io)))
            p.tlast = 0
            update!(p, length(result))
            sleep(1)
            CSI = "\e["
            print(io, "$(CSI)$(ceil(Int, ninstances/2)+1)A$(CSI)1G$(CSI)0J")
        end
    end

    result = Dict{NamedTuple{(:julia, :pkg), Tuple{VersionNumber, String}},Symbol}()
    try @sync begin
        # Printer
        @async begin
            try
                while (!isempty(jobs) || !all(==(nothing), running)) && !done
                    update_output()
                end
                println()
                stop_work()
            catch e
                stop_work()
                !isa(e, InterruptException) && rethrow(e)
            end
        end

        # Workers
        for i = 1:ninstances
            push!(all_workers, @async begin
                try
                    while !isempty(jobs) && !done
                        job = pop!(jobs)
                        times[i] = now()
                        running[i] = job
                        pkg_version, status, reason, log =
                            run_sandboxed_test(job.julia, job.pkg; kwargs...)
                        result[(julia=job.julia,pkg=job.pkg.name)] = status
                        running[i] = nothing

                        # report to the caller
                        if callback !== nothing
                            callback(job.julia, job.pkg.name, pkg_version,
                                     times[i], status, reason, log)
                        elseif log !== missing
                            mkpath(log_path(job.julia))
                            write(joinpath(log_path(job.julia), "$(job.pkg.name).log"), log)
                        end
                    end
                catch e
                    stop_work()
                    isa(e, InterruptException) || rethrow(e)
                end
            end)
        end
    end
    catch e
        e isa InterruptException || rethrow(e)
    end
    return result
end

"""
    run(julia_versions::Vector{VersionNumber}=[Base.VERSION],
        pkg_names::Vector{String}=[]]; registry=General, update_registry=true, kwargs...)

Run all tests for all packages in the registry `registry`, or only for the packages as
identified by their name in `pkgnames`, using Julia versions `julia_versions`.
The registry is first updated if `update_registry` is set to true.

Refer to `run_sandboxed_test`[@ref] and `run_sandboxed_julia`[@ref] for more possible
keyword arguments.
"""
function run(julia_versions::Vector{VersionNumber}=[Base.VERSION],
             pkg_names::Vector{String}=String[];
             registry::String=DEFAULT_REGISTRY, update_registry::Bool=true, kwargs...)
    # high-level entry-point that takes care of everything

    prepare_registry(registry; update=update_registry)

    prepare_runner()
    pkgs = read_pkgs(pkg_names)

    prepare_julia.(julia_versions)
    run(julia_versions, pkgs; kwargs...)
end


text:
using Base.Dates
using Base.Test

using MarketData

using TimeSeries

@testset "combine" begin

@testset "collapse operations" begin
    @testset "collapse squishes correctly" begin
        @test collapse(cl, week, first).values[2]    == 97.75
        @test collapse(cl, week, first).timestamp[2] == Date(2000,1,10)

        @test collapse(cl, week, first, last).values[2]    == 100.44
        @test collapse(cl, week, first, last).timestamp[2] == Date(2000,1,10)

        @test collapse(cl, month, first).values[2]    == 100.25
        @test collapse(cl, month, first).timestamp[2] == Date(2000,2,1)

        @test collapse(ohlc, week, first).values[2, :] == [102.0, 102.25, 94.75, 97.75]
        @test collapse(ohlc, week, first).timestamp[2] == Date(2000,1,10)

        @test collapse(ohlc, week, first, last).values[2, :] == [100.0, 102.25, 99.38, 100.44]
        @test collapse(ohlc, week, first, last).timestamp[2] == Date(2000,1,10)

        @test collapse(ohlc, month, first).values[2, :] == [104.0, 105.0, 100.0, 100.25]
        @test collapse(ohlc, month, first).timestamp[2] == Date(2000,2,1)
    end
end

@testset "merge works correctly" begin
    cl1  = cl[1:3]
    op1  = cl[2:4]
    aapl = tail(AAPL)
    ba   = tail(BA)

    @testset "takes colnames kwarg correctly" begin
        @test merge(cl, ohlc["High", "Low"], colnames=["a","b","c"]).colnames == ["a", "b", "c"]
        @test merge(cl, op, colnames=["a","b"]).colnames                      == ["a", "b"]
        @test_throws ErrorException merge(cl, op, colnames=["a"])
        @test_throws ErrorException merge(cl, op, colnames=["a","b","c"])

        @test merge(cl, ohlc["High", "Low"], :inner, colnames=["a","b","c"]).colnames == ["a", "b", "c"]
        @test merge(cl, op, :inner, colnames=["a","b"]).colnames == ["a", "b"]
        @test_throws ErrorException merge(cl, op, :inner, colnames=["a"])
        @test_throws ErrorException merge(cl, op, :inner, colnames=["a","b","c"])

        @test merge(cl, ohlc["High", "Low"], :left, colnames=["a","b","c"]).colnames == ["a", "b", "c"]
        @test merge(cl, op, :left, colnames=["a","b"]).colnames          == ["a", "b"]
        @test_throws ErrorException merge(cl, op, :left, colnames=["a"])
        @test_throws ErrorException merge(cl, op, :left, colnames=["a","b","c"])

        @test merge(cl, ohlc["High", "Low"], :right, colnames=["a","b","c"]).colnames == ["a", "b", "c"]
        @test merge(cl, op, :right, colnames=["a","b"]).colnames == ["a", "b"]
        @test_throws ErrorException merge(cl, op, :right, colnames=["a"])
        @test_throws ErrorException merge(cl, op, :right, colnames=["a","b","c"])

        @test merge(cl, ohlc["High", "Low"], :outer, colnames=["a","b","c"]).colnames == ["a", "b", "c"]
        @test merge(cl, op, :outer, colnames=["a","b"]).colnames == ["a", "b"]
        @test_throws ErrorException merge(cl, op, :outer, colnames=["a"])
        @test_throws ErrorException merge(cl, op, :outer, colnames=["a","b","c"])
    end

    @testset "returns correct alignment with Dates and values" begin
        @test merge(cl,op).values == merge(cl,op, :inner).values
        @test merge(cl,op).values[2,1] == cl.values[2,1]
        @test merge(cl,op).values[2,2] == op.values[2,1]
    end

    @testset "aligns with disparate sized objects" begin
        @test merge(cl, op[2:5]).values[1,1]  == cl.values[2,1]
        @test merge(cl, op[2:5]).values[1,2]  == op.values[2,1]
        @test merge(cl, op[2:5]).timestamp[1] == Date(2000,1,4)
        @test length(merge(cl, op[2:5]))      == 4

        @test length(merge(cl1, op1, :inner))     == 2
        @test merge(cl1, op1, :inner).values[2,1] == cl1.values[3,1]
        @test merge(cl1, op1, :inner).values[2,2] == op1.values[2,1]

        @test length(merge(cl1, op1, :left))     == 3
        @test merge(cl1, op1, :left).values[2,1] == cl1.values[2,1]
        @test merge(cl1, op1, :left).values[2,2] == op1.values[1,1]
        @test isnan(merge(cl1,op1, :left).values[1,2])

        @test length(merge(cl1, op1, :right))     == 3
        @test merge(cl1, op1, :right).values[2,1] == cl1.values[3,1]
        @test merge(cl1, op1, :right).values[2,2] == op1.values[2,1]
        @test isnan(merge(cl1, op1, :right).values[3,1])

        @test length(merge(cl1, op1, :outer))     == 4
        @test merge(cl1, op1, :outer).values[2,1] == cl1.values[2,1]
        @test merge(cl1, op1, :outer).values[2,2] == op1.values[1,1]
        @test isnan(merge(cl1, op1, :outer).values[1,2])
        @test isnan(merge(cl1, op1, :outer).values[4,1])
    end

    @testset "column names match the correct values" begin
        @test merge(cl, op[2:5]).colnames         == ["Close", "Open"]
        @test merge(op[2:5], cl).colnames         == ["Open", "Close"]

        @test merge(cl, op[2:5], :inner).colnames == ["Close", "Open"]
        @test merge(op[2:5], cl, :inner).colnames == ["Open", "Close"]

        @test merge(cl, op[2:5], :left).colnames  == ["Close", "Open"]
        @test merge(op[2:5], cl, :left).colnames  == ["Open", "Close"]

        @test merge(cl, op[2:5], :right).colnames == ["Close", "Open"]
        @test merge(op[2:5], cl, :right).colnames == ["Open", "Close"]

        @test merge(cl, op[2:5], :outer).colnames == ["Close", "Open"]
        @test merge(op[2:5], cl, :outer).colnames == ["Open", "Close"]
    end

    @testset "unknown method" begin
        @test_throws ArgumentError merge(cl, op, :unknown)
    end
end

@testset "vcat works correctly" begin
    @testset "concatenates time series correctly in 1D" begin
        a = TimeArray([Date(2015, 10, 01), Date(2015, 11, 01)], [15, 16], ["Number"])
        b = TimeArray([Date(2015, 12, 01)], [17], ["Number"])
        c = vcat(a, b)

        @test length(c)  == (length(a) + length(b))
        @test c.colnames == a.colnames
        @test c.colnames == b.colnames
        @test c.values   == [15, 16, 17]
    end

    @testset "concatenates time series correctly in 2D" begin
        a = TimeArray([Date(2015, 09, 01), Date(2015, 10, 01), Date(2015, 11, 01)], [[15 16]; [17 18]; [19 20]], ["Number 1", "Number 2"])
        b = TimeArray([Date(2015, 12, 01)], [18 18], ["Number 1", "Number 2"])
        c = vcat(a, b)

        @test length(c)  == length(a) + length(b)
        @test c.colnames == a.colnames
        @test c.colnames == b.colnames
        @test c.values   == [[15 16]; [17 18]; [19 20]; [18 18]]
    end

    @testset "rejects when column names do not match" begin
        a = TimeArray([Date(2015, 10, 01), Date(2015, 11, 01)], [15, 16], ["Number"])
        b = TimeArray([Date(2015, 12, 01)], [17], ["Data does not match number"])

        @test_throws ArgumentError vcat(a, b)
    end

    @testset "rejects when metas do not match" begin
        a = TimeArray([Date(2015, 10, 01), Date(2015, 11, 01)], [15, 16], ["Number"], :FirstMeta)
        b = TimeArray([Date(2015, 12, 01)], [17], ["Number"], :SecondMeta)

        @test_throws ArgumentError vcat(a, b)
    end

    @testset "rejects when dates overlap" begin
        a = TimeArray([Date(2015, 10, 01), Date(2015, 11, 01)], [15, 16], ["Number"])
        b = TimeArray([Date(2015, 11, 01)], [17], ["Number"])

        @test_throws ArgumentError vcat(a, b)
    end

    @testset "still works when dates are mixed" begin
        a = TimeArray([Date(2015, 10, 01), Date(2015, 12, 01)], [15, 17], ["Number"])
        b = TimeArray([Date(2015, 11, 01)], [16], ["Number"])
        c = vcat(a, b)

        @test length(c)  == length(a) + length(b)
        @test c.colnames == a.colnames
        @test c.colnames == b.colnames
        @test c.values   == [15, 16, 17]
        @test issorted(c.timestamp)
    end
end

@testset "map works correctly" begin
    @testset "works on both time stamps and 1D values" begin
        a = TimeArray([Date(2015, 10, 01), Date(2015, 11, 01)], [15, 16], ["Number"], :Something)
        b = map((timestamp, values) -> (timestamp + Dates.Year(1), values - 1), a)

        @test length(b)                  == length(a)
        @test b.colnames                 == a.colnames
        @test Dates.year(b.timestamp[1]) == Dates.year(a.timestamp[1]) + 1
        @test b.values[1]                == a.values[1] - 1
        @test b.meta                     == a.meta
    end

    @testset "works on both time stamps and 2D values" begin
        a = TimeArray([Date(2015, 09, 01), Date(2015, 10, 01), Date(2015, 11, 01)], [[15 16]; [17 18]; [19 20]], ["Number 1", "Number 2"])
        b = map((timestamp, values) -> (timestamp + Dates.Year(1), [values[1] + 2, values[2] - 1]), a)

        @test length(b)                  == length(a)
        @test b.colnames                 == a.colnames
        @test Dates.year(b.timestamp[1]) == Dates.year(a.timestamp[1]) + 1
        @test b.values[1, 1]             == a.values[1, 1] + 2
        @test b.values[1, 2]             == a.values[1, 2] - 1
    end

    @testset "works with order of elements that varies after modifications" begin
        a = TimeArray([Date(2015, 10, 01), Date(2015, 12, 01)], [15, 16], ["Number"])
        b = map((timestamp, values) -> (timestamp + Dates.Year((timestamp >= Date(2015, 11, 01)) ? -1 : 1), values), a)

        @test length(b) == length(a)
        @test issorted(b.timestamp)
    end
end

end  # @testset "combine"


text:
export FlexiTank

"Propellant tank with flexible size"
mutable struct FlexiTank <: Tank
    drymass::Float64
    totalmass::Float64
    propellant::Float64

    function FlexiTank(drymass::Number, totalmass::Number)
        new(drymass, totalmass, totalmass - drymass)
    end
end

# These weights are roughly based on the Electron rocket
drymass(tank::FlexiTank) = tank.drymass
totalmass(tank::FlexiTank) = tank.totalmass


text:
using ImageInTerminal, ImageCore, ColorTypes, FixedPointNumbers
using TestImages, ImageTransformations, CoordinateTransformations
using ImageMagick
using Rotations, OffsetArrays
using SparseArrays
using Test

reference_path(filename) = joinpath(dirname(@__FILE__), "reference", "$(filename).txt")

function test_reference_impl(filename, actual)
    try
        reference = replace.(readlines(reference_path(filename)), Ref("\n" => ""))
        try
            @assert reference == actual # to throw error
            @test true # to increase test counter if reached
        catch # test failed
            println("Test for \"$filename\" failed.")
            println("- REFERENCE -------------------")
            println.(reference)
            println("-------------------------------")
            println("- ACTUAL ----------------------")
            println.(actual)
            println("-------------------------------")
            if isinteractive()
                print("Replace reference with actual result? [y/n] ")
                answer = first(readline())
                if answer == 'y'
                    write(reference_path(filename), join(actual, "\n"))
                end
            else
                error("You need to run the tests interactively with 'include(\"test/runtests.jl\")' to update reference images")
            end
        end
    catch ex
        if isa(ex, SystemError) # File doesn't exist
            println("Reference file for \"$filename\" does not exist.")
            println("- NEW CONTENT -----------------")
            println.(actual)
            println("-------------------------------")
            if isinteractive()
                print("Create reference file with above content? [y/n] ")
                answer = first(readline())
                if answer == 'y'
                    write(reference_path(filename), join(actual, "\n"))
                end
            else
                error("You need to run the tests interactively with 'include(\"test/runtests.jl\")' to create new reference images")
            end
        else
            throw(ex)
        end
    end
end

# using a macro looks more consistent
macro test_reference(filename, expr)
    esc(:(test_reference_impl($filename, $expr)))
end

function ensurecolor(f, args...)
    old_color = Base.have_color
    try
        Core.eval(Base, :(have_color = true))
        return @inferred f(args...)
    finally
        Core.eval(Base, :(have_color = $old_color))
    end
end

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

# define some test images
gray_square = colorview(Gray, N0f8[0. 0.3; 0.7 1])
gray_square_alpha = colorview(GrayA, N0f8[0. 0.3; 0.7 1], N0f8[1 0.7; 0.3 0])
gray_line = colorview(Gray, N0f8[0., 0.3, 0.7, 1])
gray_line_alpha = colorview(GrayA, N0f8[0., 0.3, 0.7, 1], N0f8[1, 0.7, 0.3, 0])
rgb_line = colorview(RGB, range(0, stop=1, length=20), zeroarray, range(1, stop=0, length=20))

camera_man = testimage("camera")
lighthouse = testimage("lighthouse")
toucan = testimage("toucan")
lena = testimage("lena_color_256")

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

tests = [
    "tst_colorant2ansi.jl",
    "tst_encodeimg.jl",
    "tst_imshow.jl",
    "tst_baseshow.jl",
]

for t in tests
    @testset "$t" begin
        include(t)
    end
end


text:
module Datasets

using ..FastAI
using ..FastAI: typify

import MLUtils: MLUtils, getobs, numobs, filterobs, groupobs, mapobs,
    shuffleobs, groupobs, eachobs, ObsView
import MLDatasets
using MLUtils: mapobs, groupobs
using DataDeps
using FilePathsBase
import DataAugmentation
using FilePathsBase: filename
import FileIO
using IndirectArrays: IndirectArray
using Colors
using FixedPointNumbers
using DataFrames
using Tables
using CSV
using ShowCases
using InlineTest

include("fastaidatasets.jl")

include("batching.jl")
include("containers.jl")
include("load.jl")
include("loaders.jl")
include("recipe.jl")

include("deprecations.jl")

function __init__()
    initdatadeps()
end

export
    # primitive containers
    TableDataset,

    mapobs, eachobs, groupobs, shuffleobs, ObsView,

    # utilities
    isimagefile,
    istextfile,
    matches,
    loadfile,
    loadmask,
    pathname,
    pathparent,
    parentname,
    grandparentname,

    # datasets
    loadfolderdata,
    datasetpath,

    # recipes
    loadrecipe,
    finddatasets,
    listdatasources,
    loaddataset

end  # module


text:
export rbasis, rbspline, rbsplinu

"""
	rbasis(npts,ord,t,x,h)

Subroutine to generate a rational B-spline basis using an open knot vector.

# Arguments
- `npts::Int64` : the number of the points of the control polygon.
- `ord::Int64` : the order of the B-Spline (`deg = ord-1`).
- `t::Float64` : the parameter value of the parametric curve.
- `x::Array{Float64}` : the knot vector.
- `h::Array{Float64}` : array containing the homogeneous coordinate weighting factors

# Examples

```jldoctest
julia> x = [0.0, 0.0, 0.0, 0.0, 1.0, 2.0, 2.0, 2.0, 2.0]

julia> h = [1.0, 1.0, 0.25, 1.0, 1.0]

julia> rbasis(5, 4, 0.5, x, h)
5-element Array{Float64,1}:
 0.153846
 0.730769
 0.0769231
 0.0384615
 0.0
```
_By Giuseppe Santorelli_
"""

function rbasis(npts::Int64, ord::Int64, t::Float64, x::Array{Float64}, h::Array{Float64})::Array{Float64}

	m::Int64 = npts + ord
	r::Array{Float64} = zeros(m)

	# first order nonrational basis functions
	for i = 1 : m-1
		if (t >= x[i]) && (t < x[i+1])
			r[i] = 1
		else
			r[i] = 0
		end
	end

	# higher order nonrational basis function
	d::Float64 = 0.0
	e::Float64 = 0.0

	for k = 2 : ord

		for i = 1 : m-ord

			if r[i] != 0
				d = ((t - x[i]) * r[i]) / (x[i+k-1] - x[i])
			else
				d = 0
			end

			if r[i+1] != 0
				e = ((x[i+k] - t) * r[i+1]) / (x[i+k] - x[i+1])
			else
				e = 0
			end

			r[i] = d + e
		end
	end

	# pick up last point
	if t == x[m]
		r[npts] = 1
	end

    # calculate sum for denominator of rational basis functions
	sum::Float64 = 0.0

	for i = 1 : npts
		sum = sum + r[i] * h[i]
	end

	# form rational basis function
	if sum != 0
		for i = 1 : npts
			r[i] = r[i] * h[i] / sum
		end
	else
		r = zeros(npts)
	end

	return r[1:npts]
end

#-----------------------------------------------------------------------

"""
	rbspline(npts, ord, p1, b, h)

Subroutine to generate a rational B-spline curve using an uniform open knot vector. Returns a tuple `(P,EV)`. P is the coordinates matrix of the p1 points b-spline curve defined by the control points matrix b. EV is the 1-dimensional cellular complex.

# Arguments

- `npts::Int64` : the number of the points of the control polygon.
- `ord::Int64` : the order of the B-Spline (`deg = ord-1`).
- `p1::Int64`: number of the points to be calculated on the curve
- `b::Array{Float64}`: 2-dimensional array containing the control polygon vertices
- `h::Array{Float64}` : array containing the homogeneous coordinate weighting factors

# Examples

```jldoctest
julia> b = [0 1 5/2 4 5; 1 2 0 2 0; 0 0 0 0 0]

julia> h = [1, 1, 0.25, 1, 1]

julia> rbspline(5,3,50,b,h)[1]
3×50 Array{Float64,2}:
 0.0  0.118164  0.228377  0.331479  …  4.66852   4.77162  4.88184   5.0
 1.0  1.11652   1.22178   1.31653      0.641603  0.44733  0.233982  0.0
 0.0  0.0       0.0       0.0          0.0       0.0      0.0       0.0

julia> rbspline(5,3,50,b,h)[2]
49-element Array{Array{Int64,1},1}:
 [1, 2]
 [2, 3]
 [3, 4]
 [4, 5]
 [5, 6]
 [6, 7]
 [7, 8]
 [8, 9]
 [9, 10]
 [10, 11]
 [11, 12]
 [12, 13]
 [13, 14]
 ⋮
 [38, 39]
 [39, 40]
 [40, 41]
 [41, 42]
 [42, 43]
 [43, 44]
 [44, 45]
 [45, 46]
 [46, 47]
 [47, 48]
 [48, 49]
 [49, 50]
```

_By Giuseppe Santorelli_
"""
function rbspline(npts::Int64, ord::Int64, p1::Int64, b::Array{Float64}, h::Array{Float64})::Tuple{Array{Float64,2},Array{Array{Int64,1},1}}

	@assert length(b) == 3 * npts ("ERROR: array b not matching with parameters")
	@assert npts >= ord ("ERROR: ord > npts")

	m::Int64 = npts + ord
	P::Array{Float64} = zeros(3, p1)

	#generate the uniform knot vector
	x::Array{Float64} = knot(npts, ord)

	#calculate the points on the rational B-spline curve
	icount::Int64 = 1
	t::Float64 = 0.0
	step::Float64 = x[m] / (p1 - 1)

	for i1 = 1 : p1

		if (x[m] - t < 5e-6)
			t = x[m]
		end

		#generate the basis function
		r::Array{Float64} = rbasis(npts, ord, t, x, h)

		#generate a point
		for j = 1 : 3 #iteration over three components for every point
			for i = 1 : npts
				temp = r[i] * b[j,i]
				P[j, icount] = P[j, icount] + temp
			end
		end

		icount = icount + 1
		t = t + step
	end

	EV = Array{Int64,1}[] #1-dimensional cellular complex used for plotting the curve with Plasm package

	for i = 1 : p1 - 1
		C = Array{Int64}(2)
		C[1] = i
		C[2] = i + 1
		push!(EV, C)
	end

	return (P, EV)
end

#-----------------------------------------------------------------------

"""
	rbsplinu(npts,ord,p1,b,h)

Subroutine to generate a rational B-spline curve using an uniform periodic knot vector. Returns a tuple `(P,EV)`. P is the coordinates matrix of the p1 points b-spline curve defined by the control points matrix b. EV is the 1-dimensional cellular complex.

# Arguments

- `npts::Int64` : the number of the points of the control polygon.
- `ord::Int64` : the order of the B-Spline (`deg = ord-1`).
- `p1::Int64`: number of the points to be calculated on the curve
- `b::Array{Float64}`: 2-dimensional array containing the control polygon vertices
- `h::Array{Float64}` : array containing the homogeneous coordinate weighting factors

# Examples

```jldoctest
julia> b = [0 1 5/2 4 5; 1 2 0 2 0; 0 0 0 0 0]

julia> h = [1, 1, 0.25, 1, 1]

julia> rbspline(5,3,50,b,h)[1]
3×50 Array{Float64,2}:
 0.5  0.559434  0.615602  0.668943  …  4.33106  4.3844   4.44057  4.75
 1.5  1.55779   1.60901   1.65399      1.31653  1.22178  1.11652  0.5
 0.0  0.0       0.0       0.0          0.0      0.0      0.0      0.0

julia> rbspline(5,3,50,b,h)[2]
49-element Array{Array{Int64,1},1}:
 [1, 2]
 [2, 3]
 [3, 4]
 [4, 5]
 [5, 6]
 [6, 7]
 [7, 8]
 [8, 9]
 [9, 10]
 [10, 11]
 [11, 12]
 [12, 13]
 [13, 14]
 ⋮
 [38, 39]
 [39, 40]
 [40, 41]
 [41, 42]
 [42, 43]
 [43, 44]
 [44, 45]
 [45, 46]
 [46, 47]
 [47, 48]
 [48, 49]
 [49, 50]
```
_By Giuseppe Santorelli_
"""
function rbsplinu(npts::Int64, ord::Int64, p1::Int64, b::Array{Float64}, h::Array{Float64})::Tuple{Array{Float64,2},Array{Array{Int64,1},1}}

	@assert length(b) == 3 * npts ("ERROR: array b not matching with parameters")
	@assert npts >= ord ("ERROR: ord > npts")

	m::Int64 = npts + ord
	P::Array{Float64} = zeros(3, p1)

	#generate the uniform periodic knot vector
	x::Array{Float64} = knotu(npts, ord)

	#calculate the points on the rational B-spline
	icount::Int64 = 1
	t::Float64 = ord - 1
	step::Float64 = (npts - ord + 1) / (p1 - 1)

	for i1 = 1 : p1

	if (x[m] - t < 5e-6)
			t = x[m]
		end

		#generate the basis function
		r::Array{Float64} = rbasis(npts, ord, t, x, h)

		#generate a point
		for j = 1 : 3 #iteration over three components for every point
			for i = 1 : npts
				temp = r[i] * b[j,i]
				P[j,icount] = P[j,icount] + temp
			end
		end

		icount = icount + 1
		t = t + step
	end

	EV = Array{Int64,1}[] #1-dimensional cellular complex used for plotting the curve with Plasm package

	for i = 1 : p1 - 1
		C = Array{Int64}(2)
		C[1] = i
		C[2] = i + 1
		push!(EV, C)
	end

	return (P, EV)
end


text:
# TODO: check trees are on the same leave set

"""
    RF(tree1::T, tree2::T)::Int64 where T <:AbstractNode

Calculate the Robinson-Foulds distance between the two trees.
In its current form the function assumes the trees have identical leave sets.

Returns result of algorithm as integer.

* `tree1` : tree used to determine RF distance.

* `tree2` : tree used to determine RF distance.
"""
function RF(tree1::T, tree2::T)::Int64 where T <: GeneralNode
    initialize_tree!(tree1, height=false)
    initialize_tree!(tree2, height=false)
    r, _ = RF_int(tree1, tree2)
    return r
end

function RF_int(tree1::T, tree2::T)::Tuple{Int64, Int64} where T <: GeneralNode
    bt3 = get_bipartitions(tree1)
    bt4 = get_bipartitions(tree2)
    length(bt3)+length(bt4) - 2* length(intersect(bt3, bt4)), length(bt3)+length(bt4)
end

"""
    RF_weighted(tree1::T, tree2::T)::Float64 where T <:AbstractNode

Calculate the weighted Robinson-Foulds distance between the two trees.
The raw Robinson-Foulds distance is weighted by the maximum distance between the trees.
In its current form the function assumes the trees have identical leave sets.

Returns result of algorithm as integer.

* `tree1` : tree used to determine RF distance.

* `tree2` : tree used to determine RF distance.
"""
function RF_weighted(tree1::T, tree2::T)::Float64 where T <: GeneralNode
    initialize_tree!(tree1, height=false)
    initialize_tree!(tree2, height=false)
    rf, l1 = RF_int(tree1, tree2)
    rf / l1
end

"""
    get_bipartitions(tree::T)::Vector{Tuple} where T <:AbstractNode

Get a vector of all bipartions of `tree`.

Returns a vector containing Tuples of sets representing the bipartions.
"""
function get_bipartitions(tree::T)::Vector{Tuple} where T <:AbstractNode
    po_vect= post_order(tree)[1:end-1]
    bt = Vector{Tuple}(undef, length(po_vect))
    all_leaves = [n.name for n in get_leaves(tree)]
    for ind in eachindex(po_vect)
        elem = po_vect[ind]::T
        outset = String[]
        inset = sort([i.name for i in get_leaves(elem)])
        outset = setdiff(all_leaves, inset)
        @inbounds bt[ind] = (join(sort(inset),","), join(sort(outset),","))
    end # for
    bt
end

"""
    BHV_bounds(tree1::T, tree2::T)::Tuple{Float64, Float64} where T <:AbstractNode

This function calculates the lower and upper bounds of the geodesic in the
Billera-Holmes-Vogtman space.

Returns tuple of floats.
"""
function BHV_bounds(tree1::T, tree2::T)::Tuple{Float64, Float64} where T <:AbstractNode
    bp_t1 = get_bipartitions(tree1)
    bp_t2 = get_bipartitions(tree2)

    T1minusT2 = 0.0
    T2minusT1 = 0.0
    T1andT2 = 0.0

    for bp in bp_t1
        found1 = find_lca(tree1, String.(split(bp[1], ",")))
        if found1.root
            found1 = find_lca(tree1, String.(split(bp[2], ",")))
        end

        ind = findfirst(isequal(bp), bp_t2)
        if !isnothing(ind)
            found2 = find_lca(tree2, String.(split(bp[1], ",")))
            if found2.root
                found2 = find_lca(tree2, String.(split(bp[2], ",")))
            end
            T1andT2 += (found1.inc_length - found2.inc_length)^2
            deleteat!(bp_t2, ind)
        else
            T1minusT2 += found1.inc_length^2
        end
    end

    # only nodes which are present in the second tree are left in bp_t2
    for bp in bp_t2
        found2 = find_lca(tree2, String.(split(bp[1], ",")))
        if found2.root
            found2 = find_lca(tree2, String.(split(bp[2], ",")))
        end
        T2minusT1 += found2.inc_length^2
    end

    res_low::Float64 = T1minusT2+T2minusT1+T1andT2
    res_high::Float64 = (sqrt(T1minusT2)+sqrt(T2minusT1))^2+T1andT2

    sqrt(res_low), sqrt(res_high)
end