module MFCC

using DSP, FFTW
include("Mel.jl")
include("Filterbanks.jl")
using MFCC.Mel
using MFCC.Filterbanks

"""
    mfcc(signal::AbstractVector, window::Function = ones, windowsize = 400, overlap = 160, sr = 44100, numcep = 26)

Compute the Mel-Frequency Cepstral Coefficients of the given signal and return
an Array where each column contains the coefficients of a given frame obtained
from the stft.

**Parameters:**

`signal`: The signal on which the coefficients will be computed.

`window`: The window function to be applied on the signal by the STFT.

`windowsize`: The size of the STFT window.

`overlap`: The overlap of two adjacent windows of the STFT in samples.

`sr`: Sample rate of the signal.

`numcep`: The amount of cepstral coefficients to be extracted from the signal.

See also: [`Filterbank`](@ref)
"""
function mfcc(signal::AbstractVector,
	window::Function = ones,
	windowsize::Int = 400,
	overlap::Int = 160,
	sr::Int = 44100,
	numcep::Int = 26
)

	#Each column is a frame of stft
	sigstft = stft(signal, windowsize, overlap, window=window)

	#If the signal is real then the absolute value of it's spectrum is symmetric
	eltype(signal)<:Real && (windowsize = windowsize ÷ 2 + 1; sr ÷ 2)

	#Evaluate the periodogram of the signal
	periodgrm = map(frame->(abs2.(frame)/windowsize), sigstft)

	#Create the filterbank
	windows, windowstarts = mffilters(numcep, windowsize, sr)
	melfilterbank = Filterbank(windowsize, windows[1:end], windowstarts[1:end])

	#Apply the filterbank and sum the result of each filter
	filteredsignal = mapslices(frame->melfilterbank*frame, periodgrm, dims = 1)

	#Take the Type-III DCT of the log for each coeficient
	return mapslices(frame->idct(log.(frame)), filteredsignal, dims=1)

end

export mfcc

end # module


@testset "multipole" begin
    #let's check that rotation works properly
    k0 = 0.1
    kin = 100k0
    λ0 = 2π/k0
    θ_i = 0.0

    N = 1000
    P = 5
    shapes = [squircle(0.2λ0, N)]
    centers = [-1.5λ0 0; 1.5λ0 0]
    Ns = size(centers,1)
    φs = [0.0; 0.23]
    ids = ones(Int, Ns)
    sp = ScatteringProblem(shapes, ids, centers, φs)
    points = λ0*[5 0; -2 3]

    verify_min_distance(shapes, centers, ids, points)
    β1, σ1 = solve_particle_scattering(k0, kin, P, sp, PlaneWave(θ_i);
        get_inner = true, verbose = false)
    Ez1 = scattered_field_multipole(k0, β1, centers, points)

    #now rotate everything by some angle and compare
    θ_r = 1.23#2π*rand()
    #notation is transposed due to structure of centers
    shapes2 = [squircle(0.2λ0, N); CircleParams(15)] #no effect,for extra code coverage
    centers2 = centers*[cos(θ_r) sin(θ_r); -sin(θ_r) cos(θ_r)]
    φs2 = θ_r .+ φs
    sp2 = ScatteringProblem(shapes2, ids, centers2, φs2)
    points2 = points*[cos(θ_r) sin(θ_r); -sin(θ_r) cos(θ_r)]

    β2, σ2 = solve_particle_scattering(k0, kin, P, sp2, PlaneWave(θ_i + θ_r);
        get_inner = true, verbose = true)
    Ez2 = scattered_field_multipole(k0, β2, centers2, points2)

    @test Ez1 ≈ Ez2
    β2_r = copy(β2)
    for ic = 1:Ns, p = -P:P
        β2_r[(ic-1)*(2P+1) + p + P + 1] *= exp(1.0im*θ_r*p)
    end
    @test β1 ≈ β2_r
    @test σ1 ≈ σ2
end

@testset "circle" begin
    k0 = 1
    kin = 0.1
    R = 0.2
    P = 10
    S, gamma = ParticleScattering.circleScatteringMatrix(k0, kin, R, P; gamma = true)
    center = [0 2R]
    points = center .+ [0 0.5R; 0.1R -0.2R]
    E1 = ParticleScattering.innerFieldCircle(kin, gamma, center[1,:], points)
    E2 = ParticleScattering.innerFieldCircle(kin, gamma, center[1,:], points[1,:])
    @test E1[1] == E2
end


solve(i::BudgetedBipartiteMatchingInstance{T}, ::LagrangianAlgorithm, ε; kwargs...) where T = matching_hungarian_budgeted_lagrangian_search(i, ε; kwargs...)
solve(i::BudgetedBipartiteMatchingInstance{T}, ::LagrangianRefinementAlgorithm; kwargs...) where T = matching_hungarian_budgeted_lagrangian_refinement(i; kwargs...)
solve(i::BudgetedBipartiteMatchingInstance{T}, ::IteratedLagrangianRefinementAlgorithm; kwargs...) where T = matching_hungarian_budgeted_lagrangian_approx_half(i; kwargs...)

approximation_term(::BudgetedBipartiteMatchingInstance{T}, ::LagrangianAlgorithm) where T = NaN
approximation_ratio(::BudgetedBipartiteMatchingInstance{T}, ::LagrangianAlgorithm) where T = NaN

approximation_term(i::BudgetedBipartiteMatchingInstance{T}, ::LagrangianRefinementAlgorithm) where T = maximum(values(i.matching.rewards))
approximation_ratio(::BudgetedBipartiteMatchingInstance{T}, ::LagrangianRefinementAlgorithm) where T = NaN

approximation_term(::BudgetedBipartiteMatchingInstance{T}, ::IteratedLagrangianRefinementAlgorithm) where T = NaN
approximation_ratio(::BudgetedBipartiteMatchingInstance{T}, ::IteratedLagrangianRefinementAlgorithm) where T = 0.5

struct BudgetedBipartiteMatchingLagrangianSolution{T, U} <: BudgetedBipartiteMatchingSolution{T, U}
  # Used to store important temporary results from solving the Lagrangian dual.
  instance::BudgetedBipartiteMatchingInstance{T, U}
  solution::Vector{Edge{T}}
  λ::Float64 # Optimum dual multiplier.
  value::Float64 # Optimum value of the dual problem (i.e. with penalised constraint).
  λmax::Float64 # No dual value higher than this is useful (i.e. they all yield the same solution).
end

struct SimpleBudgetedBipartiteMatchingSolution{T, U} <: BudgetedBipartiteMatchingSolution{T, U}
  instance::BudgetedBipartiteMatchingInstance{T, U}
  solution::Vector{Edge{T}}
  value::Float64

  function SimpleBudgetedBipartiteMatchingSolution(instance::BudgetedBipartiteMatchingInstance{T, U}, solution::Vector{Edge{T}}) where {T, U}
    if length(solution) > 0
      return new{T, U}(instance, solution, _budgeted_bipartite_matching_compute_value(instance, solution))
    else
      return new{T, U}(instance, solution, -Inf)
    end
  end

  function SimpleBudgetedBipartiteMatchingSolution(instance::BudgetedBipartiteMatchingInstance{T, U}) where {T, U}
    # No feasible solution.
    return new{T, U}(instance, edgetype(instance.matching.graph)[], -Inf)
  end
end

function matching_hungarian_budgeted_lagrangian(i::BudgetedBipartiteMatchingInstance{T, U}, λ::Float64) where {T, U}
  # Solve the subproblem for one value of the dual multiplier λ:
  #     l(λ) = \max_{x matching} (rewards + λ weights) x - λ budget.
  bmi_rewards = Dict{Edge{T}, Float64}(e => i.matching.reward[e] + λ * i.weight[e] for e in keys(i.matching.reward))
  bmi = BipartiteMatchingInstance(i.matching, bmi_rewards)
  bmi_sol = matching_hungarian(bmi)
  return _budgeted_bipartite_matching_compute_value(bmi, bmi_sol.solution) - λ * i.budget, bmi_sol.solution
end

function matching_hungarian_budgeted_lagrangian_search(i::BudgetedBipartiteMatchingInstance{T, U}, ε::Float64) where {T, U}
  # Approximately solve the problem \min_{l ≥ 0} l(λ), where
  #     l(λ) = \max_{x smatching} (rewards + λ weights) x - λ budget.
  # This problem is the Lagrangian dual of the budgeted maximum spanning-tree problem:
  #     \max_{x matching} rewards x  s.t.  weights x >= budget.
  # This algorithm provides no guarantee on the optimality of the solution.

  # TODO: generalise this function? Highly similar to the one for spanning trees.

  # Initial set of values for λ. The optimum is guaranteed to be contained in this interval.
  weights_norm_inf = maximum(values(i.weight)) # Maximum weight.
  m = min(i.matching.n_left, i.matching.n_right) # Maximum number of items in a solution. Easy to compute for a matching!
  λmax = Float64(weights_norm_inf * m + 1)

  λlow = 0.0
  λhigh = λmax

  # Perform a golden-ratio search.
  while (λhigh - λlow) > ε
    λmidlow = λhigh - (λhigh - λlow) / MathConstants.φ
    λmidhigh = λlow + (λhigh - λlow) / MathConstants.φ
    vmidlow, _ = matching_hungarian_budgeted_lagrangian(i, λmidlow)
    vmidhigh, _ = matching_hungarian_budgeted_lagrangian(i, λmidhigh)

    if vmidlow < vmidhigh
      λhigh = λmidhigh
      vhigh = vmidhigh
    else
      λlow = λmidlow
      vlow = vmidlow
    end
  end

  vlow, bmlow = matching_hungarian_budgeted_lagrangian(i, λlow)
  vhigh, bmhigh = matching_hungarian_budgeted_lagrangian(i, λhigh)
  if vlow < vhigh
    return BudgetedBipartiteMatchingLagrangianSolution(i, bmlow, λlow, vlow, λmax)
  else
    return BudgetedBipartiteMatchingLagrangianSolution(i, bmhigh, λhigh, vhigh, λmax)
  end
end

function matching_hungarian_budgeted_lagrangian_refinement(i::BudgetedBipartiteMatchingInstance{T, U};
                                                           ε::Float64=1.0e-3, ζ⁻::Float64=0.2, ζ⁺::Float64=5.0,
                                                           stalling⁻::Float64=1.0e-5) where {T, U}
  # Approximately solve the following problem:
  #     \max_{x matching} rewards x  s.t.  weights x >= budget
  # This algorithm provides an additive approximation to this problem. If x* is the optimum solution and x~ the one
  # returned by this algorithm,
  #     weights x* >= budget   and   weights x~ >= budget                 (the returned solution is feasible)
  #     rewards x~ >= rewards x* - 2 \max{e edge} reward[e]               (additive approximation)
  # This is based on http://people.idsia.ch/~grandoni/Pubblicazioni/BBGS08ipco.pdf, Lemma 4.

  # Check assumptions.
  if ζ⁻ >= 1.0
    error("ζ⁻ must be strictly less than 1.0: the dual multiplier λ is multiplied by ζ⁻ to reach an infeasible solution by less penalising the budget constraint.")
  end
  if ζ⁺ <= 1.0
    error("ζ⁺ must be strictly greater than 1.0: the dual multiplier λ is multiplied by ζ⁺ to reach a feasible solution by penalising more the budget constraint.")
  end

  # Ensure the problem is feasible by only considering the budget constraint.
  feasible_rewards = Dict{Edge{T}, Float64}(e => i.weight[e] for e in keys(i.weight))
  feasible_instance = BipartiteMatchingInstance(i.matching, feasible_rewards)
  feasible_solution = matching_hungarian(feasible_instance)
  if _budgeted_bipartite_matching_compute_value(feasible_instance, feasible_solution.solution) < i.budget
    # By maximising the left-hand side of the budget constraint, impossible to reach the target budget. No solution!
    return SimpleBudgetedBipartiteMatchingSolution(i)
  end

  # Solve the Lagrangian relaxation to optimality.
  lagrangian = matching_hungarian_budgeted_lagrangian_search(i, ε)
  λ0, v0, bm0 = lagrangian.λ, lagrangian.value, lagrangian.solution
  λmax = lagrangian.λmax
  b0 = _budgeted_bipartite_matching_compute_weight(i, bm0) # Budget consumption of this first solution.

  # If already respecting the budget constraint exactly, done!
  if b0 == i.budget
    return SimpleBudgetedBipartiteMatchingSolution(i, bm0)
  end

  # Find two solutions: one above the budget x⁺ (i.e. respecting the constraint), the other not x⁻.
  x⁺, x⁻ = nothing, nothing

  λi = λ0
  if b0 > i.budget
    x⁺ = bm0
    @assert _budgeted_bipartite_matching_compute_weight(i, x⁺) > i.budget

    stalling = false
    while true
      # Penalise less the constraint: it should no more be satisfied.
      λi *= ζ⁻
      _, sti = matching_hungarian_budgeted_lagrangian(i, λi)
      if _budgeted_bipartite_matching_compute_weight(i, sti) < i.budget
        x⁻ = sti
        break
      end

      # Is the process stalling?
      if λi <= stalling⁻
        stalling = true
        break
      end
    end

    # Specific handling of stallings.
    if stalling # First test: don't penalise the constraint at all.
      _, sti = matching_hungarian_budgeted_lagrangian(i, 0.0)
      new_budget = _budgeted_bipartite_matching_compute_weight(i, sti)
      if new_budget < i.budget
        x⁻ = sti
        stalling = false
      end
    end

    if stalling # Second test: minimise the left-hand side of the budget constraint, in hope of finding a feasible solution.
      # This process is highly similar to the computation of feasible_solution, but with a reverse objective function.
      infeasible_rewards = Dict{Edge{T}, Float64}(e => - i.weight[e] for e in keys(i.weight))
      infeasible_solution = matching_hungarian(BipartiteMatchingInstance(i.matching.graph, infeasible_rewards)).solution

      if _budgeted_bipartite_matching_compute_weight(i, infeasible_solution) < i.budget
        x⁻ = infeasible_solution
        stalling = false
      end
    end

    if stalling # Third: decide there is no solution strictly below the budget. No refinement is possible.
      # As x⁺ is feasible, return it.
      return SimpleBudgetedBipartiteMatchingSolution(i, x⁺)
    end
  else
    x⁻ = bm0
    @assert _budgeted_bipartite_matching_compute_weight(i, x⁻) < i.budget

    while true
      # Penalise more the constraint: it should become satisfied at some point.
      λi *= ζ⁺
      _, sti = matching_hungarian_budgeted_lagrangian(i, λi)
      if _budgeted_bipartite_matching_compute_weight(i, sti) >= i.budget
        x⁺ = sti
        break
      end

      # Is the process stalling? If so, reuse feasible_solution, which is guaranteed to be feasible.
      if λi >= λmax
        x⁺ = feasible_solution.solution
        break
      end
    end
  end

  # Normalise the solutions: the edges point from the left part to the right part of the bipartite graph.
  sort_edge(e::Edge{T}) where T = (i.matching.partition[src(e)] == 1) ? e : reverse(e)
  x⁺ = [sort_edge(e) for e in x⁺]
  x⁻ = [sort_edge(e) for e in x⁻]

  already_seen_solutions = Set{Set{Edge{T}}}()
  push!(already_seen_solutions, Set(x⁺))
  push!(already_seen_solutions, Set(x⁻))

  # Iterative refinement. Stop as soon as there is a difference of at most one edge between the two solutions.
  while _solution_symmetric_difference_size(x⁺, x⁻) > 4
    # Enforce the loop invariant.
    @assert x⁺ !== nothing
    @assert x⁻ !== nothing
    @assert _budgeted_bipartite_matching_compute_weight(i, x⁺) >= i.budget # Feasible.
    @assert _budgeted_bipartite_matching_compute_weight(i, x⁻) < i.budget # Infeasible.

    # Switch elements from one solution to another: use the XOR between the two solutions as an augmenting path/cycle.
    only_in_x⁺, only_in_x⁻ = _solution_symmetric_difference(x⁺, x⁻)
    xor = vcat(only_in_x⁺, only_in_x⁻)

    current_edge = first(xor)
    current_vertex = src(current_edge)
    all_vertices = Set{T}(current_vertex)
    xor = xor[2:end]
    new_x = copy(x⁺)

    while true
      # Exit conditions:
      # - If the XOR is a path: consume all edges; at some point, will no more
      #   be possible to find the next one
      # - Otherwise, once a node is met twice, break.

      # If the current edge is in the solution, remove it. Otherwise, add it.
      if current_edge in new_x
        filter!(e -> e != current_edge, new_x)
      else
        push!(new_x, current_edge)
      end

      # Go to the next edge, following a cycle or a path.
      if current_vertex == src(current_edge)
        current_vertex = dst(current_edge)
      else
        current_vertex = src(current_edge)
      end
      current_edge_idx = findfirst(e -> src(e) == current_vertex || dst(e) == current_vertex, xor)
      if isnothing(current_edge_idx)
        break
      end
      current_edge = xor[current_edge_idx]
      deleteat!(xor, current_edge_idx)

      # Check for a cycle.
      if current_vertex in all_vertices
        break
      end
      push!(all_vertices, current_vertex)
    end

    # If this procedure cannot propose a new solution, it has stalled and will not recover.
    if Set(new_x) in already_seen_solutions
      break
    else
      push!(already_seen_solutions, Set(new_x))
    end

    # Replace one of the two solutions, depending on whether this solution is feasible (x⁺) or not (x⁻).
    if _budgeted_bipartite_matching_compute_weight(i, new_x) >= i.budget
      x⁺ = new_x
    else
      x⁻ = new_x
    end
  end

  # Done!
  return SimpleBudgetedBipartiteMatchingSolution(i, x⁺)
end

function _edge_any_end_match(e::Edge{T}, t::Edge{T}) where T
  return src(e) == src(t) || dst(e) == dst(t) || src(e) == dst(t) || dst(e) == src(t)
end

function matching_hungarian_budgeted_lagrangian_approx_half(i::BudgetedBipartiteMatchingInstance{T, U}; kwargs...) where {T, U}
  # Approximately solve the following problem:
  #     \max_{x matching} rewards x  s.t.  weights x >= budget
  # This algorithm provides a multiplicative approximation to this problem. If x* is the optimum solution and x~ the one
  # returned by this algorithm,
  #     weights x* >= budget   and   weights x~ >= budget                 (the returned solution is feasible)
  #     rewards x~ >= (1 - ε) rewards x*                                  (multiplicative approximation)
  # The parameter ε is not tuneable, but rather fixed to 1/2.

  # This is based on http://people.idsia.ch/~grandoni/Pubblicazioni/BBGS08ipco.pdf, Theorem 1.

  # If there are too few vertices, not much to do. The smallest side of the bipartite graph must have at least 4 vertices,
  # so that 4 of them can be fixed.
  if i.matching.n_left <= 4 || i.matching.n_right <= 4
    # TODO: Should rather perform an exhaustive exploration.
    return matching_hungarian_budgeted_lagrangian_refinement(i; kwargs...)
  end

  # For each combination of four distinct edges, force these four edges to be part of the solution and discard all edges with a higher value.
  # If selecting twice the same edge: no need to go further.
  # If any end of two selected edges are the same: this cannot lead to a feasible solution.
  # Both tests are implemented in _edge_any_end_match.
  best_sol = nothing
  for e1 in edges(i.matching.graph)
    for e2 in edges(i.matching.graph)
      if _edge_any_end_match(e1, e2)
        continue
      end

      for e3 in edges(i.matching.graph)
        if _edge_any_end_match(e1, e3) || _edge_any_end_match(e2, e3)
          continue
        end

        for e4 in edges(i.matching.graph)
          if _edge_any_end_match(e1, e4) || _edge_any_end_match(e2, e4) || _edge_any_end_match(e3, e4)
            continue
          end

          # Filter out the edges that have a higher value than any of these two edges. Give a very large reward to them both.
          cutoff = min(i.matching.reward[e1], i.matching.reward[e2], i.matching.reward[e3], i.matching.reward[e4])
          reward = copy(i.matching.reward)
          filter!(kv -> kv[2] <= cutoff, reward)
          filter!(kv -> ! _edge_any_end_match(kv[1], e1), reward)
          filter!(kv -> ! _edge_any_end_match(kv[1], e2), reward)
          filter!(kv -> ! _edge_any_end_match(kv[1], e3), reward)
          filter!(kv -> ! _edge_any_end_match(kv[1], e4), reward)
          reward[e1] = reward[e2] = reward[e3] = reward[e4] = prevfloat(Inf)

          if length(keys(reward)) == 4 # Nothing left? Skip.
            continue
          end

          graph = SimpleGraph(nv(i.matching.graph))
          for e in keys(reward)
            add_edge!(graph, e)
          end

          weight = Dict(e => i.weight[e] for e in keys(reward))

          # Solve this subproblem.
          bbmi = BudgetedBipartiteMatchingInstance(graph, reward, weight, i.budget)
          sol = matching_hungarian_budgeted_lagrangian_refinement(bbmi; kwargs...)

          # This subproblem is infeasible. Maybe it's because the overall problem is infeasible or just because too many
          # edges were removed.
          if length(sol.solution) == 0
            continue
          end

          # Impossible to have a feasible solution with these two edges, probably because of the budget constraint.
          # It's very unlikely that this happens, due to the checks already performed when looping.
          if ! (e1 in sol.solution) || ! (e2 in sol.solution) || ! (e3 in sol.solution) || ! (e4 in sol.solution)
            continue
          end

          # Only keep the best solution.
          if best_sol == nothing || _budgeted_bipartite_matching_compute_value(i, sol.solution) > _budgeted_bipartite_matching_compute_value(i, best_sol.solution)
            # sol's instance is the one used internally for the subproblems.
            best_sol = SimpleBudgetedBipartiteMatchingSolution(i, sol.solution)
          end
        end
      end
    end
  end

  if best_sol !== nothing
    return best_sol
  else
    return SimpleBudgetedBipartiteMatchingSolution(i)
  end
end


module RandomExtensions

export Combine, Uniform, Normal, Exponential, CloseOpen, Rand

import Random: Sampler, rand, rand!

using Random
using Random: GLOBAL_RNG, SamplerTrivial, SamplerSimple, SamplerTag, Repetition

using SparseArrays: sprand, sprandn

include("distributions.jl")
include("sampling.jl")
include("containers.jl")
include("iteration.jl")

## updated rand docstring (TODO: replace Base's one)

"""
    rand([rng=GLOBAL_RNG], [S], [C...]) # RandomExtensions

Pick a random element or collection of random elements from the set of values specified by `S`;
`S` can be

* an indexable collection (for example `1:n` or `['x','y','z']`),
* an `AbstractDict` or `AbstractSet` object,
* a string (considered as a collection of characters), or
* a type: the set of values to pick from is then equivalent to `typemin(S):typemax(S)` for
  integers (this is not applicable to [`BigInt`](@ref)), and to ``[0, 1)`` for floating
  point numbers;
* a `Distribution` object, e.g. `Normal()` for a normal distribution (like `randn()`),
  or `CloseOpen(10.0, 20.0)` for uniform `Float64` numbers in the range ``[10.0, 20.0)``;
* a `Combine` object, which can be either `Combine(Pair, S1, S2)` or `Combine(Complex, S1, S2)`,
  where `S1` and `S2` are one of the specifications above; `Pair` or `Complex` can optionally be
  given as concrete types, e.g. `Combine(ComplexF64, 1:3, Int)` to generate `ComplexF64` instead
  of `Complex{Int}`.

`S` usually defaults to [`Float64`](@ref).

If `C...` is not specified, `rand` produces a scalar. Otherwise, `C...` can be:

* a set of integers, or a tuple of `Int`, which specify the dimensions of an `Array` to generate;
* `(p::AbstractFloat, m::Integer, [n::Integer])`, which produces a sparse array of dimensions `(m, n)`,
  in which the probability of any element being nonzero is independently given by `p`
* `(String, [n=8])`, which produces a random `String` of length `n`; the generated string consists of `Char`
  taken from a predefined set like `randstring`, and can be specified with the `S` parameter.
* `(Dict, n)`, which produces a `Dict` of length `n`; `S` must then specify the type of its elements,
  e.g. `Combine(Pair, Int, 2:3)`;
* `(Set, n)`, which produces a `Set` of length `n`;
* `(BitArray, dims...)`, which produces a `BitArray` with the specified dimensions.

# Examples
```julia-repl
julia> rand(Int, 2)
2-element Array{Int64,1}:
 1339893410598768192
 1575814717733606317

julia> rand(MersenneTwister(0), Dict(1=>2, 3=>4))
1=>2

julia> rand("abc", String, 12)
"bccaacaabaac"

julia> rand(1:10, Set, 3)
Set([3, 8, 6])
```

!!! note
    The complexity of `rand(rng, s::Union{AbstractDict,AbstractSet})`
    is linear in the length of `s`, unless an optimized method with
    constant complexity is available, which is the case for `Dict`,
    `Set` and `BitSet`. For more than a few calls, use `rand(rng,
    collect(s))` instead, or either `rand(rng, Dict(s))` or `rand(rng,
    Set(s))` as appropriate.
"""
rand

end # module


##############
# SumOfTerms #
##############

struct SumOfLocalTerms{N, T1<:NTuple{N,ScalarFunction},
                            T2<:NTuple{N,LocalOperator}} <: LocalOperator
    coefficients::T1
    operators::T2
end

operators(op::SumOfLocalTerms) = op.operators
coefficients(op::SumOfLocalTerms) = op.coefficients

Base.:*(α::ScalarFunction, op::LocalOperator) = SumOfLocalTerms((α,), (op,))
Base.:*(op::LocalOperator, α::ScalarFunction) = *(α, op)

Base.:*(α::ScalarFunction, op::SumOfLocalTerms) =
    SumOfLocalTerms(α .* op.coefficients, op.operators)
Base.:*(op::SumOfLocalTerms, α::ScalarFunction) = *(α, op)

Base.:-(op::LocalOperator) = SumOfLocalTerms((-1,), (op,))
Base.:-(op::SumOfLocalTerms) = SumOfLocalTerms(.-(op.coefficients), op.operators)

Base.:+(op1::LocalOperator, op2::LocalOperator) = SumOfLocalTerms((1,1), (op1, op2))
Base.:-(op1::LocalOperator, op2::LocalOperator) = SumOfLocalTerms((1,-1), (op1, op2))
Base.:+(op1::SumOfLocalTerms, op2::LocalOperator) =
    SumOfLocalTerms((op1.coefficients..., 1), (op1.operators..., op2))
Base.:-(op1::SumOfLocalTerms, op2::LocalOperator) =
    SumOfLocalTerms((op1.coefficients..., -1), (op1.operators..., op2))
Base.:+(op1::LocalOperator, op2::SumOfLocalTerms) =
    SumOfLocalTerms((1, op2.coefficients...), (op1, op2.operators...))
Base.:-(op1::LocalOperator, op2::SumOfLocalTerms) = +(op1, -(op2))
Base.:+(op1::SumOfLocalTerms, op2::SumOfLocalTerms) =
    SumOfLocalTerms((op1.coefficients..., op2.coefficients...),
                    (op1.operators..., op2.operators...))
Base.:-(op1::SumOfLocalTerms, op2::SumOfLocalTerms) = +(op1, -(op2))

Base.:*(op1::SumOfLocalTerms, op2::LocalOperator) =
    SumOfLocalTerms(op1.coefficients, op1.operators .* (op2,))
Base.:*(op1::LocalOperator, op2::SumOfLocalTerms) =
    SumOfLocalTerms(op2.coefficients, (op1,) .* op2.operators)

Base.:*(op1::SumOfLocalTerms{1}, op2::SumOfLocalTerms) =
    op1.coefficients[1] * (op1.operators[1] * op2)

function Base.:*(op1::SumOfLocalTerms, op2::SumOfLocalTerms)
    op1.coefficients[1] * (op1.operators[1] * op2) +
    SumOfLocalTerms(Base.tail(op1.coefficients), Base.tail(op1.operators)) * op2
end

Base.adjoint(op::SumOfLocalTerms) =
    SumOfLocalTerms(adjoint.(op.coefficients), adjoint.(op.operators))


"""
 Tests for https://arxiv.org/abs/1303.5103 with test problems as in the paper.
 DRI1, RI1, RI3, RI5, RI6, RDI1WM, RDI2WM, RDI3WM, RDI4WM
"""

import Statistics # for mean values of trajectories
import LinearAlgebra # for the normn
using StochasticDiffEq
using Test
using Random
#using DiffEqGPU

#using Plots

function generate_weak_solutions(prob, alg, dts, numtraj; ensemblealg=EnsembleThreads())
  sols = []
  for i in 1:length(dts)
    sol = solve(prob,alg;ensemblealg=ensemblealg,dt=dts[i],save_start=false,save_everystep=false,weak_timeseries_errors=false,weak_dense_errors=false,trajectories=Int(numtraj))
    println(i)
    push!(sols,sol)
  end
  return sols
end

function prob_func(prob, i, repeat)
    remake(prob,seed=seeds[i])
end

"""
 Test Scalar SDEs (oop)
"""

@info "Scalar oop noise"

numtraj = Int(2e5) # in the paper they use 1e9
u₀ = 0.0
f(u,p,t) = 1//2*u+sqrt(u^2+1)
g(u,p,t) = sqrt(u^2+1)
dts = 1 .//2 .^(4:-1:1)
tspan = (0.0,2.0) # 2.0 in paper

h1(z) = z^3-6*z^2+8*z
#analytical_sol(t) = E(f(X(t))) = E(h1(arsinh(X(t))) = t^3-3*t^2+2*t
#analytical_sol(2) = 0 and analytical_sol(1)=0

seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

prob = SDEProblem(f,g,u₀,tspan)
ensemble_prob = EnsembleProblem(prob;
        output_func = (sol,i) -> (h1(asinh(sol[end])),false),
        prob_func = prob_func
        )
_solutions = @time generate_weak_solutions(ensemble_prob, DRI1(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

#convergence_plot = plot(dts, errors, xaxis=:log, yaxis=:log)
#savefig(convergence_plot, "DRI1-CPU-final-"*string(numtraj)*".pdf")
println("DRI1:", m)

numtraj = Int(5e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

_solutions = @time generate_weak_solutions(ensemble_prob, RI1(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RI1:", m)

numtraj = Int(5e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

_solutions = @time generate_weak_solutions(ensemble_prob, RI3(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RI3:", m)

numtraj = Int(5e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

_solutions = @time generate_weak_solutions(ensemble_prob, RI5(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RI5:", m)

numtraj = Int(7e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

_solutions = @time generate_weak_solutions(ensemble_prob, RI6(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RI6:", m)

numtraj = Int(1e3)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

_solutions = @time generate_weak_solutions(ensemble_prob, RDI1WM(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-1) < 0.3

println("RDI1WM:", m)

numtraj = Int(7e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

_solutions = @time generate_weak_solutions(ensemble_prob, RDI2WM(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RDI2WM:", m)

_solutions = @time generate_weak_solutions(ensemble_prob, RDI3WM(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

numtraj = Int(3e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

println("RDI3WM:", m)

_solutions = @time generate_weak_solutions(ensemble_prob, RDI4WM(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RDI4WM:", m)

"""
 Test Scalar SDEs (iip)
"""

@info "Scalar iip noise"

u₀ = [0.0]
f1!(du,u,p,t) = @.(du = 1//2*u+sqrt(u^2 +1))
g1!(du,u,p,t) = @.(du = sqrt(u^2 +1))
dts = 1 .//2 .^(4:-1:1)
tspan = (0.0,2.0)

h1(z) = z^3-6*z^2+8*z

prob = SDEProblem(f1!,g1!,u₀,tspan)
ensemble_prob = EnsembleProblem(prob;
        output_func = (sol,i) -> (h1(asinh(sol[end][1])),false),
        prob_func = prob_func
        )

numtraj = Int(2e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

_solutions = @time generate_weak_solutions(ensemble_prob, DRI1(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("DRI1:", m)

numtraj = Int(5e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

_solutions = @time generate_weak_solutions(ensemble_prob, RI1(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RI1:", m)

numtraj = Int(5e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

_solutions = @time generate_weak_solutions(ensemble_prob, RI3(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RI3:", m)

numtraj = Int(5e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

_solutions = @time generate_weak_solutions(ensemble_prob, RI5(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RI5:", m)

numtraj = Int(7e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

_solutions = @time generate_weak_solutions(ensemble_prob, RI6(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RI6:", m)

numtraj = Int(1e3)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

_solutions = @time generate_weak_solutions(ensemble_prob, RDI1WM(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-1) < 0.3

println("RDI1WM:", m)

numtraj = Int(7e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

_solutions = @time generate_weak_solutions(ensemble_prob, RDI2WM(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RDI2WM:", m)

numtraj = Int(3e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

_solutions = @time generate_weak_solutions(ensemble_prob, RDI3WM(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RDI3WM:", m)

_solutions = @time generate_weak_solutions(ensemble_prob, RDI4WM(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RDI4WM:", m)

"""
 Test non-commutative noise SDEs (iip)
"""

@info "Non-commutative noise"

u₀ = [1.0,1.0]
function f2!(du,u,p,t)
  du[1] = -273//512*u[1]
  du[2] = -1//160*u[1]-(-785//512+sqrt(2)/8)*u[2]
end
function g2!(du,u,p,t)
  du[1,1] = 1//4*u[1]
  du[1,2] = 1//16*u[1]
  du[2,1] = (1-2*sqrt(2))/4*u[1]
  du[2,2] = 1//10*u[1]+1//16*u[2]
end
dts = 1 .//2 .^(3:-1:0)
tspan = (0.0,3.0)

h2(z) = z^2 # but apply it only to u[1]

prob = SDEProblem(f2!,g2!,u₀,tspan,noise_rate_prototype=zeros(2,2))
ensemble_prob = EnsembleProblem(prob;
        output_func = (sol,i) -> (h2(sol[end][1]),false),
        prob_func = prob_func
        )

numtraj = Int(1e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

_solutions = @time generate_weak_solutions(ensemble_prob, DRI1(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)-exp(-3.0)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("DRI1:", m)

numtraj = Int(2e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)
_solutions = @time generate_weak_solutions(ensemble_prob, RI1(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)-exp(-3.0)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RI1:", m)

numtraj = Int(2e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)
_solutions = @time generate_weak_solutions(ensemble_prob, RI3(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)-exp(-3.0)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RI3:", m)

numtraj = Int(2e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)
_solutions = @time generate_weak_solutions(ensemble_prob, RI5(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)-exp(-3.0)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RI5:", m)

numtraj = Int(2e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)
_solutions = @time generate_weak_solutions(ensemble_prob, RI6(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)-exp(-3.0)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RI6:", m)

numtraj = Int(1e3)
seed = 10
Random.seed!(seed)
seeds = rand(UInt, numtraj)

_solutions = @time generate_weak_solutions(ensemble_prob, RDI1WM(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)-exp(-3.0)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-1) < 0.3

println("RDI1WM:", m)

numtraj = Int(1e5)
seed = 10
Random.seed!(seed)
seeds = rand(UInt, numtraj)

_solutions = @time generate_weak_solutions(ensemble_prob, RDI2WM(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)-exp(-3.0)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RDI2WM:", m)

_solutions = @time generate_weak_solutions(ensemble_prob, RDI3WM(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)-exp(-3.0)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RDI3WM:", m)

_solutions = @time generate_weak_solutions(ensemble_prob, RDI4WM(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)-exp(-3.0)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RDI4WM:", m)

"""
 Test Diagonal noise SDEs (iip), SIAM Journal on Numerical Analysis, 47 (2009), pp. 1713–1738
"""

@info "Diagonal noise"

u₀ = [0.1,0.1]
function f3!(du,u,p,t)
  du[1] = 3//2*u[1]
  du[2] = 3//2*u[2]
end
function g3!(du,u,p,t)
  du[1] = 1//10*u[1]
  du[2] = 1//10*u[1]
end
dts = 1 .//2 .^(3:-1:0)
tspan = (0.0,1.0)

h3(z) = z^2 # == 1//10**exp(3//2*t) if h3(z) = z and  == 1//100**exp(301//100*t) if h3(z) = z^2 )

prob = SDEProblem(f3!,g3!,u₀,tspan)
ensemble_prob = EnsembleProblem(prob;
        output_func = (sol,i) -> (h3(sol[end][1]),false),
        prob_func = prob_func
        )

numtraj = Int(5e4)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

_solutions = @time generate_weak_solutions(ensemble_prob, DRI1(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)-1//100*exp(301//100)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("DRI1:", m)

_solutions = @time generate_weak_solutions(ensemble_prob, RI1(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)-1//100*exp(301//100)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RI1:", m)

_solutions = @time generate_weak_solutions(ensemble_prob, RI3(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)-1//100*exp(301//100)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RI3:", m)

_solutions = @time generate_weak_solutions(ensemble_prob, RI5(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)-1//100*exp(301//100)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RI5:", m)

numtraj = Int(1e5)
seed = 70
Random.seed!(seed)
seeds = rand(UInt, numtraj)

_solutions = @time generate_weak_solutions(ensemble_prob, RI6(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)-1//100*exp(301//100)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.45

println("RI6:", m)

numtraj = Int(1e3)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

_solutions = @time generate_weak_solutions(ensemble_prob, RDI1WM(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)-1//100*exp(301//100)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-1) < 0.45

println("RDI1WM:", m)

numtraj = Int(1e5)
seed = 70
Random.seed!(seed)
seeds = rand(UInt, numtraj)

_solutions = @time generate_weak_solutions(ensemble_prob, RDI2WM(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)-1//100*exp(301//100)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.45

println("RDI2WM:", m)

numtraj = Int(5e4)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

_solutions = @time generate_weak_solutions(ensemble_prob, RDI3WM(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)-1//100*exp(301//100)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RDI3WM:", m)

_solutions = @time generate_weak_solutions(ensemble_prob, RDI4WM(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)-1//100*exp(301//100)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RDI4WM:", m)


# this doesn't actually put anything into the type form
# It transforms any number into a bitstype that can hold that number

#function make_number_from_tuple( n::Tuple ; base = 1000 )
#        return foldl( (a,b)-> a*base+b, big(0), n )
#end
function make_number_from_tuple( n::Tuple ; base = 1000 )
        # BigInts to ensure known return type
        fn(a::BigInt,b) = BigInt( a*base+b )
        return foldl( fn, big(0), n )
end

function make_tuple_from_number( n; last_digits = tuple(), base =1000 )
        if -base < n < base
                return ( oftype(base,n),last_digits...)
        end
        next_n, next_digit = fldmod(n,base)
        return make_tuple_from_number( next_n; last_digits = (oftype( base, next_digit ),last_digits...), base = base )
end

# for testing:
function tuple_number_round_trip( n ; base =1000 )
        t = make_tuple_from_number( n; base=base )
        n2 = make_number_from_tuple( t; base=base )
        @assert( n == n2 )
        return n2
end

function to_type_constant( n::Integer )
        if     Int == Int32 && typemin(Int32)  <= n <= typemax(Int32)
                # only for machines where Int is Int32
                return Int32(n)
        elseif typemin(Int64)  <= n <= typemax(Int64)
                return Int64(n)
        elseif typemin(Int128) <= n <= typemax(Int128)
                return Int128(n)
        else
                return make_tuple_from_number(n )
        end
end

function from_type_constant( n::Integer )
        return BigInt(n)
end

function from_type_constant( n::Tuple )
        return make_number_from_tuple( n )
end


export Martin

struct Martin <: LakeModel end

@with_kw_noshow mutable struct MartinParameters <: LakeParameters
    @deftype Float64
    nutrients = 0.7 # current nutrient level
    ib = 2e-5 # g⋅m⁻²⋅day⁻¹ immigration rate of bream (9e-3 g⋅m⁻²⋅year⁻¹)
    ip = 2e-5 # g⋅m⁻²⋅day⁻¹ immigration rate of pike (9e-3 g⋅m⁻²⋅year⁻¹)
    r = 7.5e-3 # day⁻¹ maximum growth rate of bream (2.74 year⁻¹)
    H₁ = 0.5 # half saturation constant
    H₂ = 11 # % half saturation constant
    H₃ = 20 # g⋅m⁻² half saturation constant
    H₄ = 15 # g⋅m⁻² half saturation constant
    cb = 7.5e-5 # m⁻²⋅g⁻¹⋅day⁻¹ intraspecific competition constant for bream (0.0274 m⁻²⋅g⁻¹⋅year⁻¹)
    cp = 2.75e-4 # m⁻²⋅g⁻¹⋅day⁻¹ intraspecific competition constant for pike (0.1 m⁻²⋅g⁻¹⋅year⁻¹)
    prmax = 5e-2 # day⁻¹ maximum predation rate of pike (18.25 year⁻¹)
    ce = 0.1 # pike food conversion efficiency to growth
    mp = 2.25e-3 # day⁻¹ mortality rate of pike (0.82 year⁻¹)
    rv = 7e-3 # day⁻¹ vetegation growth rate
    cv = 6e-5 # m² intraspecific competition for vetetation
    mv = 7e-3 # day⁻¹ mortality of vegetation
    pv = 0.0 # day planting rate
    tb = 0.0 # day trawling rate
end

function lake_dynamics!(du, u, p::MartinParameters, t)
    B, P, V = u # g⋅m⁻² bream/pike/vegetation density

    FR = B^2 / (B^2 + p.H₄^2) # fuctional response of pike
    # Wrapped Cauchy distribution, simplifed to a yearly growth cycle between
    # March & July.
    # Expanded: γ*(1-ρ^2)/(1+ρ^2-2ρ*cos(2π*(t-μ)/τ))
    # NOTE: Having some issues with this, it's probably not necessary at all anyway.
    #plant = p.pv > 0.0 ? -p.pv / (p.pv + 19.4872 * cos(2π / 365.0 * (t - 122.0))) : 0.0

    du[1] =
        dB =
            p.ib + p.r * (p.nutrients / (p.nutrients + p.H₁)) * B - p.tb * B - p.cb * B^2 -
            p.prmax * FR * P
    du[2] = dP = p.ip + p.ce * p.prmax * FR * P * (V / (V + p.H₂)) - p.mp * P - p.cp * P^2
    du[3] = dV = p.rv * V + p.pv * V - p.cv * V^2 - p.mv * (V * B^2 / (p.H₃^2 + B^2))
end

"""
    lake_initial_state(nutrients, bream, pike, vegetation, Martin)

Provides initial conditions of the lake with specified values.
"""
function lake_initial_state(
    nutrients::Float64,
    bream::Float64,
    pike::Float64,
    vegetation::Float64,
    ::Type{Martin};
    kwargs...,
)
    ([bream, pike, vegetation], MartinParameters(; nutrients = nutrients, kwargs...))
end

"""
    lake_initial_state(nutrients, <:LakeDefinition, Martin)

A special override that returns interpolated results from precaclulated `Clear` and
`Turbid` bifurcation results. Not as precise as the preset conditions, since the
numbers do not come directly from the analytical solution, but close.
"""
function lake_initial_state(
    nutrients::Float64,
    ::Type{L},
    ::Type{Martin};
    kwargs...,
) where {L<:LakeDefinition}
    initial_state = interpolated_position(nutrients, L)
    (initial_state, MartinParameters(; nutrients = nutrients, kwargs...))
end

LakeParameters(::Type{Martin}, nutr; kwargs...) =
    MartinParameters(; nutrients = nutr, kwargs...)

preset_conditions(::Type{Clear}, ::Type{Martin}) = (0.7, [20.5172, 1.7865, 56.8443])
preset_conditions(::Type{Turbid}, ::Type{Martin}) = (2.5, [83.0128, 0.0414705, 6.40048])
preset_conditions(::Type{X1}, ::Type{Martin}) = (2.2, [36.730, 2.87725, 26.6798]) #deep unstable
preset_conditions(::Type{X2}, ::Type{Martin}) = (1.05, [59.0606, 0.819124, 12.0023]) #unstable
preset_conditions(::Type{X3}, ::Type{Martin}) = (1.05, [64.0559, 0.374008, 10.3631]) #stable

preset_conditions(::Type{S1}, ::Type{Martin}) = (2.0, [79.597, 0.050, 6.928]) #turbid transition, stable
preset_conditions(::Type{S2}, ::Type{Martin}) = (3.5, [85.436, 0.0373, 6.061]) #turbid eutrophic, stable
preset_conditions(::Type{S3}, ::Type{Martin}) = (3.5, [20.0, 3.5, 50.0]) #clear eutrophic, unstable
preset_conditions(::Type{T1}, ::Type{Martin}) = (0.9, [21.647, 2.050, 53.726]) #clear oligotrophic, high pike concentration, stable
preset_conditions(::Type{T2}, ::Type{Martin}) = (2.0, [26.901, 2.809, 41.530]) #clear transition, stable
preset_conditions(::Type{T3}, ::Type{Martin}) = (3.5, [20.0, 3.5, 50.0]) #clear eutrophic unstable

function Base.show(io::IO, ::MIME"text/plain", p::MartinParameters)
    println(io, "Parameters for lake dynamics (Martin) model):")
    println(io, "Nutrient level: $(p.nutrients)")
    println(io, "Immigration rate (g⋅m⁻²⋅day⁻¹) for bream: $(p.ib), pike: $(p.ip)")
    println(
        io,
        "Growth rate (day⁻¹) for bream: $(p.r), Predation rate (day⁻¹) of pike: $(p.prmax)",
    )
    println(
        io,
        "Half sturation constants: H₁ $(p.H₁), H₂ $(p.H₂)%, H₃ $(p.H₃) g⋅m⁻², H₄ $(p.H₄) g⋅m⁻²",
    )
    println(
        io,
        "Intraspecific competition constant (g⋅m⁻²⋅day⁻¹) for bream: $(p.cb), pike: $(p.cp)",
    )
    println(
        io,
        "Pike food conversion efficiency to growth: $(p.ce), Mortalitiy rate (day⁻¹): $(p.mp)",
    )
    println(
        io,
        "Vegetation rates. Growth: $(p.rv) day⁻¹, mortality: $(p.mv) day⁻¹, competition: $(p.cv) m²",
    )
end

function interpolated_position(nutrients::Float64, ::Type{Turbid})
    @assert 1.04 <= nutrients <= 4.0
    [turbid_b(nutrients), turbid_p(nutrients), turbid_v(nutrients)]
end

function interpolated_position(nutrients::Float64, ::Type{Clear})
    @assert 0.0 <= nutrients <= 2.49
    [clear_b(nutrients), clear_p(nutrients), clear_v(nutrients)]
end

# Higher order interpolations were a bit jumpy on the ends. Since we account for this already in the
# grid, there's no real point doing much more than linear interpolation in these data.
const turbid_n = [1.04, 1.06, 1.08, 1.1, 1.12, 1.14, 1.16, 1.18, 1.2, 1.22, 1.24, 1.26, 1.28, 1.3, 1.4, 1.5, 1.6, 1.7, 1.8, 1.9, 2.0, 2.1, 2.2, 2.3, 2.4, 2.8, 2.9, 3.0, 3.1, 3.2, 3.3, 3.4, 3.5, 3.6, 3.7, 3.8, 3.9, 4.0]
const turbid_b = LinearInterpolation([62.8433, 64.7811, 65.812, 66.5967154617846, 67.256, 67.8374, 68.3651, 68.8529, 69.3096380295164, 69.7411, 70.1515, 70.5439, 70.9205, 71.2830815128419, 72.927620226924, 74.3575735324184, 75.6261770482514, 76.7652279581881, 77.7965051730699, 78.7361588920137, 79.5967630382622, 80.3884228650353, 81.1194401014282, 81.7967455181591, 82.4261974404837, 84.5557605975222, 85.0087274579361, 85.4355006715746, 85.838305423903, 86.2191189576847, 86.5797047356448, 86.9216409781937, 87.2463446598708, 87.5550918073417, 87.849034761985, 88.1292169339137, 88.3965854697749, 88.6520021756033], turbid_n)
const turbid_p = LinearInterpolation([0.467635, 0.324727, 0.264457, 0.226438572671889, 0.199488, 0.179158, 0.163257, 0.150262, 0.139566501222224, 0.130573, 0.122817, 0.116197, 0.110491, 0.105282876267212, 0.086680414110511, 0.075007876580953, 0.066965656681847, 0.06106181970993, 0.056627123629523, 0.053158110844447, 0.05023878194799, 0.048013362468672, 0.046125185620789, 0.044221413526291, 0.042868037308363, 0.038674849804663, 0.038099615841367, 0.037304090525497, 0.036805912699434, 0.036048380454724, 0.035464861024496, 0.034753892992486, 0.034831406585964, 0.034185163193823, 0.03418612062571, 0.033622695610838, 0.032871258345589, 0.032617613483335], turbid_n)
const turbid_v = LinearInterpolation([10.7297, 10.1524, 9.86357, 9.6515996372934, 9.47861, 9.32976, 9.1976, 9.07784, 8.96776090326038, 8.86554, 8.76989, 8.67984, 8.59468, 8.51383001117352, 8.16073696425393, 7.87085225645183, 7.62612135404186, 7.4157645933494, 7.2325673056776, 7.07136843964007, 6.92831478514543, 6.80044251320062, 6.68542064458636, 6.58138361478245, 6.48681686627753, 6.18128018305646, 6.11902788139653, 6.0612162952174, 6.00738633014564, 5.95714022580627, 5.91013159862657, 5.86605737616413, 5.82465121270593, 5.78567807610134, 5.74892976796485, 5.71422119441633, 5.68138724304635, 5.6502801549636], turbid_n)

const clear_n = [0.0, 0.04, 0.06, 0.08, 0.1, 0.12, 0.14, 0.16, 0.18, 0.2, 0.22, 0.24, 0.26, 0.28, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 1.1, 1.2, 1.3, 1.4, 1.5, 1.6, 1.7, 1.8, 1.9, 2.0, 2.1, 2.2, 2.3, 2.4, 2.42, 2.44, 2.46, 2.48, 2.49]
const clear_b = LinearInterpolation([0.0, 7.19822, 10.2447, 12.7489, 14.3417899783703, 15.1403, 15.6098, 15.9564, 16.2464, 16.5043680973213, 16.7416, 16.9638, 17.1743, 17.3752, 17.5681143534832, 18.4409411744765, 19.2038264444943, 19.8895014439516, 20.5171998487337, 21.0999740586095, 21.6473927297074, 22.1668544296341, 22.6643378784209, 23.1448798085877, 23.6129093438797, 24.0725068928104, 24.5276307687052, 24.9823468293504, 25.4411004976985, 25.9090895977944, 26.3928443866286, 26.9012392076907, 27.4474789812928, 28.0536139494919, 28.7632478716902, 29.6944205304346, 29.9376, 30.2195, 30.5678, 31.0733, 31.6065], clear_n)
const clear_p = LinearInterpolation([0.0111003, 0.0142037, 0.023943, 0.0485642, 0.105562347388399, 0.190455, 0.282928, 0.374336, 0.462298, 0.546230121593651, 0.626104, 0.702078, 0.774371, 0.843219, 0.908850307777514, 1.19586574516796, 1.42898170571722, 1.6226301580547, 1.78649633470377, 1.92730359110776, 2.04986983505402, 2.15774343788958, 2.2535990505578, 2.33949577013192, 2.41704615336661, 2.48753542023639, 2.55200022043938, 2.6112914060315, 2.66611434179448, 2.71706494128768, 2.76465434572004, 2.80932964683044, 2.85149627675782, 2.89154743171501, 2.92990689414219, 2.96718342309789, 2.97461, 2.98208, 2.98966, 2.99752, 3.00196], clear_n)
const clear_v = LinearInterpolation([109.96, 103.287, 92.4176, 82.9578, 77.0475035887902, 74.1649, 72.5014, 71.2896, 70.2871, 69.4037351989893, 68.5992, 67.8522, 67.1506, 66.4863, 65.8539683061925, 63.0572420785788, 60.7016894081954, 58.6565110822716, 56.8443323127489, 55.2131743708244, 53.7257083300374, 52.3539004349192, 51.0759395776379, 49.8743098980489, 48.7344963127633, 47.6440564653745, 46.5919019868462, 45.5676811402653, 44.5611694769215, 43.5615585328629, 42.5564678917703, 41.5303333021048, 40.4613483013786, 39.3146251420862, 38.023089573663, 36.4083114241203, 36.0011, 35.536, 34.9723, 34.1743, 33.3579], clear_n)


s = open("input.txt") do file
    readlines(file);
end

function RunProgramUntilRepeat(program)
    seen = Set()

    i = 1;
    acc = 0;
    while !in(i, seen) && i <= length(program)
        push!(seen, i);
        parts = split(program[i], " ");

        cmd = parts[1];
        num = parse(Int, parts[2]);

        if cmd == "nop"
            i += 1;
        elseif cmd == "acc"
            acc += num;
            i += 1;
        elseif cmd == "jmp"
            i += num
        end
    end

    if i == length(program) + 1
        return acc
    end

    return -999999;
end

function FixProgram(program)
    for i = 1:length(program)
        if occursin("nop", program[i])
            program[i] = replace(program[i], "nop"=>"jmp");
            val = RunProgramUntilRepeat(program);

            if val != -999999
                return val
            end
            program[i] = replace(program[i], "jmp"=>"nop");

        elseif occursin("jmp", program[i])
            program[i] = replace(program[i], "jmp"=>"nop");
            val = RunProgramUntilRepeat(program);

            if val != -999999
                return val
            end
            program[i] = replace(program[i], "nop"=>"jmp");
        end
    end
end

FixProgram(s)


using BasisFunctions, BasisFunctions.Test, DoubleFloats

using Test
⊕ = BasisFunctions.:⊕
⊗ = BasisFunctions.:⊗

types = (Float64,LargeFloat)

function test_operators(T)
    @testset "identity operator" begin
        test_identity_operator(T) end

    @testset "diagonal operators" begin
        test_diagonal_operators(T) end

    @testset "multidiagonal operators" begin
        test_multidiagonal_operators(T) end

    @testset "test invertible operators" begin
        test_invertible_operators(T) end

    @testset "noninvertible operators" begin
        test_noninvertible_operators(T) end

    @testset "tensor operators" begin
        test_tensor_operators(T) end

    @testset "circulant operator" begin
        test_circulant_operator(T) end

    @testset "banded operators" begin
        test_banded_operator(T) end

    @testset "sparse operator" begin
        test_sparse_operator(T) end
end

function test_generic_operators(T)
    b1 = Fourier{T}(3)
    b2 = ChebyshevT{T}(4)
    b3 = ChebyshevT{T}(3)
    b4 = Legendre{T}(3)

    cartlist = [CartesianIndex(1,1),CartesianIndex(2,3)]
    operators = [
        ["Scaling operator", ScalingOperator(b1, T(2))],
        ["Zero operator", BF.ZeroOperator(b1, b2)],
        ["Diagonal operator", DiagonalOperator(b2, b2, rand(T, length(b2)))],
        ["Block operator", differentiation(b1⊕b2)],
        ["Coefficient scaling operator", BF.CoefficientScalingOperator(b1, 1, T(2))],
        ["Index restriction operator", IndexRestriction(b2, b1, 1:3) ],
        ["Index restriction operator (2)", IndexRestriction(b1⊗b2, cartlist) ],
        ["Index extension operator", IndexExtension(b1, b2, 1:3) ],
        ["Index extension operator (2)", IndexExtension(b1⊗b2, cartlist) ],
        ["Linearization operator", BasisFunctions.LinearizationOperator(b1⊗b2) ],
        ["DelinearizationOperator operator", BasisFunctions.DelinearizationOperator(b1⊗b2) ],
    ]

    for ops in operators
        @testset "$(ops[1])" begin
            test_generic_operator_interface(ops[2], T)
        end
    end
end

function test_tensor_operators(T)
    m1 = 3
    n1 = 4
    m2 = 10
    n2 = 24
    A1 = ArrayOperator(rand(T,m1,n1))
    A2 = ArrayOperator(rand(T,m2,n2))

    A_tp = TensorProductOperator(A1, A2)
    b = rand(T, n1, n2)
    c_tp = zeros(T, m1, m2)
    apply!(A_tp, c_tp, b)

    # Apply the operators manually, row by row and then column by column
    intermediate = zeros(T, m1, n2)
    result1 = zeros(T, m1)
    for i = 1:n2
        apply!(A1, result1, b[:,i])
        intermediate[:,i] = result1
    end
    c_manual = similar(c_tp)
    result2 = zeros(T, m2)
    for i = 1:m1
        apply!(A2, result2, reshape(intermediate[i,:], n2))
        c_manual[i,:] = result2
    end

    @test sum(abs.(c_tp-c_manual)) + 1 ≈ 1
end

function test_diagonal_operators(T)
    for SRC in (Fourier{T}(10), ChebyshevT{T}(11))
        operators = (BF.CoefficientScalingOperator(SRC, 3, rand(coefficienttype(SRC))),
            BF.AlternatingSignOperator(SRC), IdentityOperator(SRC),
            ScalingOperator(SRC, rand(coefficienttype(SRC))), ScalingOperator(SRC,3),
            DiagonalOperator(SRC, rand(coefficienttype(SRC), size(SRC))))
           # PseudoDiagonalOperator(SRC, map(coefficienttype(SRC), rand(size(SRC)))))
        for Op in operators
            m = matrix(Op)
            # Test in-place
            if isinplace(Op)
                coef_src = rand(src(Op))
                coef_dest_m = m * coef_src
                coef_dest = apply!(Op, coef_src)
                @test sum(abs.(coef_dest-coef_dest_m)) + 1 ≈ 1
            end
            # Test out-of-place
            coef_src = rand(src(Op))
            coef_dest = apply(Op, coef_src)
            coef_dest_m = m * coef_src
            @test sum(abs.(coef_dest-coef_dest_m)) + 1 ≈ 1
            # Test inverse
            I = inv(Op)
            @test abs(sum(abs.(I*(Op*coef_src)-coef_src))) + 1 ≈ 1
            @test abs(sum(abs.((I*Op)*coef_src-coef_src))) + 1 ≈ 1
            # Test transposes
            @test abs(sum(abs.(m'*coef_src-Op'*coef_src))) + 1 ≈ 1
            # Test Sum
            @test abs(sum(abs.((Op+Op)*coef_src-2*(Op*coef_src)))) + 1 ≈ 1
            # Test Equivalence to diagonal operator
            @test abs(sum(abs.(Op*coef_src-diag(Op).*coef_src))) + 1 ≈ 1
            # Make sure diagonality is retained
            @test isdiag(2*Op)
            @test isdiag(Op')
            @test isdiag(inv(Op))
            @test isdiag(Op*Op)
            # Make sure in_place is retained
            if isinplace(Op)
                @test isinplace(2*Op)
                @test isinplace(Op')
                @test isinplace(inv(Op))
                @test isinplace(Op*Op)
            end
        end
    end
end

function test_multidiagonal_operators(T)
    MSet = Fourier{T}(10)⊕ChebyshevT{T}(11)
    operators = (BF.CoefficientScalingOperator(MSet, 3, rand(coefficienttype(MSet))),
        BF.AlternatingSignOperator(MSet), IdentityOperator(MSet),
        ScalingOperator(MSet,2.0+2.0im), ScalingOperator(MSet, 3),
        DiagonalOperator(MSet, rand(coefficienttype(MSet),length(MSet))))
    for Op in operators
        # Test in-place
        coef_src = zeros(MSet)
        for i in eachindex(coef_src)
            coef_src[i] = rand(eltype(coef_src))
        end
        m = matrix(Op)
        coef_dest_m = m * linearize_coefficients(MSet,coef_src)
        coef_dest = linearize_coefficients(MSet, apply!(Op, coef_src))
        @test sum(abs.(coef_dest-coef_dest_m)) + 1 ≈ 1
        # Test out-of-place
        coef_src = zeros(MSet)
        for i in eachindex(coef_src)
            coef_src[i] = rand() * one(eltype(coef_src))
        end
        coef_dest = linearize_coefficients(MSet, apply(Op, coef_src))
        coef_dest_m = m * linearize_coefficients(MSet, coef_src)
        @test sum(abs.(coef_dest-coef_dest_m)) + 1 ≈ 1

        # Test inverse
        I = inv(Op)
        @test I*(Op*coef_src)≈coef_src
        @test (I*Op)*coef_src≈coef_src
        # Test transposes
        @test m'*linearize_coefficients(MSet,coef_src)≈linearize_coefficients(MSet,Op'*coef_src)
        # Test Sum
        @test (Op+Op)*coef_src≈2*(Op*coef_src)
        # Test Equivalence to diagonal operator
        @test linearize_coefficients(MSet,Op*coef_src)≈ diag(Op).*linearize_coefficients(MSet,coef_src)
        # Make sure diagonality is retained
        @test isdiag(2*Op) && isinplace(2*Op)
        @test isdiag(Op') && isinplace(Op')
        @test isdiag(inv(Op)) && isinplace(inv(Op))
        @test isdiag(Op*Op) && isinplace(Op*Op)
    end
end

function test_sparse_operator(ELT)
    S = SparseMatrixOperator(ArrayOperator(rand(ELT,4,4)))
    test_generic_operator_interface(S, ELT)
end

function test_banded_operator(ELT)
    a = [ELT(1),ELT(2),ELT(3)]
    H = HorizontalBandedOperator(Fourier{ELT}(6), Fourier{ELT}(3),a,3,2)
    test_generic_operator_interface(H, ELT)
    h = matrix(H)
    e = zeros(ELT,6,1)
    for i in 1:3
        e .= 0
        I = (2+(i-1)*3) .+ (1:3)
        for (j,k) in enumerate(mod.(I .- 1,6) .+ 1)
            e[k] = a[j]
        end
        @test e ≈ h[i,:]
    end

    V = VerticalBandedOperator(Fourier{ELT}(3), Fourier{ELT}(6),a,3,2)
    test_generic_operator_interface(V, ELT)
    v = matrix(V)
    e = zeros(ELT,6,1)
    for i in 1:3
        e .= 0
        I = (2+(i-1)*3) .+ (1:3)
        for (j,k) in enumerate(mod.(I .- 1,6) .+ 1)
            e[k] = a[j]
        end
        @test e ≈ v[:,i]
    end
end

function test_circulant_operator(ELT)
    n = 20
    for T in (ELT, complex(ELT))
        e1 = zeros(T,n); e1[1] = 1
        c = rand(T, n)
        C = CirculantOperator(c)
        m = matrix(C)
        for i in 1:n
            @test m[:,i] ≈ circshift(c,i-1)
        end
        invC = inv(C)
        @test BasisFunctions.eigvals(invC).*BasisFunctions.eigvals(C) ≈ ones(T,n)
        @test m'*e1 ≈ C'*e1

        sumC = invC+C
        @test (typeof(sumC) <:BasisFunctions.CirculantOperator)
        @test sumC*e1 ≈ m[:,1] + matrix(invC)[:,1]

        subC = invC-C
        @test (typeof(subC) <:BasisFunctions.CirculantOperator)
        @test subC*e1 ≈ -m[:,1] + matrix(invC)[:,1]

        mulC = invC*C
        @test (typeof(sumC) <:BasisFunctions.CirculantOperator)
        @test mulC*e1 ≈ (m*matrix(invC))[:,1]

        mulC = 2*C
        @test (typeof(sumC) <:BasisFunctions.CirculantOperator)
        @test mulC*e1 ≈ 2*C*e1

        mulC = C*2
        @test (typeof(sumC) <:BasisFunctions.CirculantOperator)
        @test mulC*e1 ≈ 2*C*e1
    end
end

function test_identity_operator(T)
    b = Fourier{T}(10)
    I = IdentityOperator(b)
    @test I isa IdentityOperator
    @test eltype(I) == Complex{T}
    @test src(I) == b
    I = IdentityOperator(Fourier{T}(10), ChebyshevT{T}(10))
    @test src(I) == Fourier{T}(10)
    @test dest(I) == ChebyshevT{T}(10)
end

function test_invertible_operators(T)
    for SRC in (Fourier{T}(10), ChebyshevT{T}(10))
        operators = (MultiplicationOperator(SRC, SRC, rand(coefficienttype(SRC),length(SRC),length(SRC))),
            CirculantOperator(rand(T, 10)),
            CirculantOperator(rand(complex(T), 10)))
        for Op in operators
            m = matrix(Op)
            # Test out-of-place
            coef_src = rand(src(Op))
            coef_dest = apply(Op, coef_src)
            coef_dest_m = m * coef_src
            @test sum(abs.(coef_dest-coef_dest_m)) + 1 ≈ 1
            # Test inverse
            I = inv(Op)
            @test abs(sum(abs.(I*(Op*coef_src)-coef_src))) + 1 ≈ 1
            @test abs(sum(abs.((I*Op)*coef_src-coef_src))) + 1 ≈ 1
            # Test transposes
            @test abs(sum(abs.(m'*coef_src-Op'*coef_src))) + 1 ≈ 1
            # Test Sum
            @test abs(sum(abs.((Op+Op)*coef_src-2*(Op*coef_src)))) + 1 ≈ 1
        end
    end
end

# FunctionOperator is currently not tested, since we cannot assume it is a linear operator.
function test_noninvertible_operators(T)
    for SRC in (Fourier{T}(10), ChebyshevT{T}(11))
        operators = (MultiplicationOperator(SRC, resize(SRC,length(SRC)+2), rand(coefficienttype(SRC),length(SRC)+2,length(SRC))),
            BF.ZeroOperator(SRC))
        for Op in operators
            m = matrix(Op)
            # Test out-of-place
            coef_src = rand(src(Op))
            coef_dest = apply(Op, coef_src)
            coef_dest_m = m * coef_src
            @test sum(abs.(coef_dest-coef_dest_m)) + 1 ≈ 1
            # Test transposes
            @test (sum(abs.(m'-matrix(Op')))) + 1 ≈ 1
            # Test Sum
            @test (sum(abs.((Op+Op)*coef_src-2*(Op*coef_src)))) + 1 ≈ 1
        end
    end
end

for T in types
    delimit(string(T))
    test_operators(T)
    test_generic_operators(T)
end


using Test
using UnPack
using MOONS.Grids
const G = Grids

@testset "Coordinates init" begin
  for FT in (Float32, Float64)
    a = FT(0.0)
    b = FT(1.0)
    n = 10
    c = Coordinates(a,b,n)
    @test c.c.Δh[1] ≈ FT(1/10)
    @test c.n.Δh[1] ≈ FT(1/10)
    @test h_extreme(c.n, Min()) ≈ FT(0)
    @test h_extreme(c.n, Max()) ≈ FT(1)
    sprint(show, c)
    @test c.c == coords(c, Center1D())
    @test c.n == coords(c, Vertex1D())
    # FT==Float64 && @show c
  end
end

@testset "Helper funcs" begin
    @test n_hat(Min()) == -1
    @test binary(Min()) == 0
    @test ghost_vec(Min()) == (1, 0, 0)
    @test ghost_dual(Min()) == (1, 0)

    @test n_hat(Max()) == 1
    @test binary(Max()) == 1
    @test ghost_vec(Max()) == (0, 0, 1)
    @test ghost_dual(Max()) == (0, 1)
    @test orthogs(1) == (2, 3)
    @test orthogs(2) == (3, 1)
    @test orthogs(3) == (1, 2)
end

@testset "Stretch funcs" begin
    FT = Float64
    Re = FT(100)
    Ha = FT(200)
    hmin = FT(0)
    hmax = FT(1)
    δ = hmax*99/100
    @test all(G.ReynoldsBL(Re, FT[0, 0, 0], FT[1, 1, 1]) .≈ [1.0540925533894598, 1.0540925533894598, 1.0540925533894598])
    @test all(G.Re_Ha_BL(Re, Ha, FT[0, 0, 0], FT[1, 1, 1]) .≈ [1.002509414234171, 1.002509414234171, 1.002509414234171])
    @test G.Re_Ha_BL_1D(Re, Ha, hmin, hmax) ≈ 1.002509414234171
    @test G.Reynolds_BL_1D(Re, hmin, hmax) ≈ 1.0540925533894598
    @test G.Hartmann_BL_1D(Ha, hmin, hmax) ≈ 1.002509414234171
    @test G.RobertsBL1D(δ, hmax) ≈ FT(1000)
    @test G.β_Δh_big(hmin, hmax, 100, FT(0.02)) ≈ 1.0443565558255983
    @test G.β_Δh_both(hmin, hmax, 100, FT(0.001)) ≈ 1.0211975072167816
    @test G.β_Δh_small(hmin, hmax, 100, FT(0.001)) ≈ 1.0218301816802464

    Δh = FT(0.1)
    @test all(uniform_left(hmin, Δh, 3) .≈ [-0.3, -0.2, -0.1, 0.0])
    @test all(uniform_right(hmin, Δh, 3) .≈ [0.0, 0.1, 0.2, 0.3])
end

@testset "Grid init" begin
  for FT in (Float32, Float64)

    a = FT(0.0)
    b = FT(1.0)
    n = 10
    c = Coordinates(a,b,n)
    @test length(c.c.h) == c.c.s
    @test length(c.n.h) == c.n.s
    @test c.n.h[1] < a
    @test c.n.h[end] > b
    @test c.n.h[2] ≈ a
    @test c.n.h[end-1] ≈ b
    @test c.c.h[1] < a
    @test c.c.h[end] > b

    a = FT(0.0)
    b = FT(1.0)
    n = 10
    β = FT(1.1)
    c1 = Coordinates(a,b,n; warpfun=Roberts_both, args=(β,))
    c2 = Coordinates(a,b,n; warpfun=Roberts_both, args=(β,))
    c3 = Coordinates(a,b,n; warpfun=Roberts_both, args=(β,))
    grid = Grid((c1,c2,c3))

    yc = FT(0.5)
    τ = FT(10)
    c1 = Coordinates(a,b,n; warpfun=cluster, args=(yc,τ))
    c2 = Coordinates(a,b,n; warpfun=cluster, args=(yc,τ))
    c3 = Coordinates(a,b,n; warpfun=cluster, args=(yc,τ))
    grid = Grid((c1,c2,c3))

    c1 = Coordinates(a,b,n; warpfun=Roberts_right, args=(β,))
    c2 = Coordinates(a,b,n; warpfun=Roberts_left, args=(β,))
    c3 = Coordinates(a,b,n; warpfun=Roberts_both, args=(β,))
    grid = Grid((c1,c2,c3))
    sprint(show, grid)
    # FT==Float64 && @show grid
    @unpack xn, yn, zn, xc, yc, zc = all_coords(grid)
    @test xn == grid.c[1].n.h
    @test yn == grid.c[2].n.h
    @test zn == grid.c[3].n.h
    @test xc == grid.c[1].c.h
    @test yc == grid.c[2].c.h
    @test zc == grid.c[3].c.h

    # Test failures
    if FT === Float32
        @test_throws ArgumentError("Coordinates are not finite.") Coordinates(a,b,n; warpfun=Roberts_left, args=(FT(1.000000000000001),))
    end

    @test over_points(c1, All(), Vertex1D()) == 1:13
    @test over_points(c1, Interior(), Vertex1D()) == 3:11
    @test over_points(c1, Boundary(), Vertex1D()) == 2:12
    @test over_points(c1, Ghost(), Vertex1D()) == (1:1, 13:13)

    @test over_points(c1, All(), Center1D()) == 1:12
    @test over_points(c1, Interior(), Center1D()) == 2:11
    @test over_points(c1, Ghost(), Center1D()) == (1:1, 12:12)
  end
end

@testset "Stretched parameter matching" begin
  for FT in (Float32, Float64)
    a = FT(0.0)
    b = FT(1.0)
    n = 10
    β = FT(1.1)
    c = Coordinates(a,b,n; warpfun=Roberts_both, args=(β,))
    β = G.β_Δh_small(a, b, n, c.n.Δh[1])
    @test β ≈ 1.1284008533596548
    @test length(c.c.h) == c.c.s
    @test length(c.n.h) == c.n.s
    @test c.n.h[1] < a
    @test c.n.h[end] > b
    @test c.n.h[2] ≈ a
    @test c.n.h[end-1] ≈ b
    @test c.c.h[1] < a
    @test c.c.h[end] > b
  end
end

@testset "Field-like indexing" begin
    FT = Float64
    a = FT(0.0)
    b = FT(1.0)
    n = 10
    β = FT(1.1)
    grid = Cube(a,b,n)
    @test all(grid.c[1].c.h .≈ [grid.cent.x[i,1,1] for i in 1:grid.c[1].c.s])
    @test all(grid.c[2].c.h .≈ [grid.cent.y[1,j,1] for j in 1:grid.c[2].c.s])
    @test all(grid.c[3].c.h .≈ [grid.cent.z[1,1,k] for k in 1:grid.c[3].c.s])

    @test all(grid.c[1].n.h .≈ [grid.corn.x[i,1,1] for i in 1:grid.c[1].n.s])
    @test all(grid.c[2].n.h .≈ [grid.corn.y[1,j,1] for j in 1:grid.c[2].n.s])
    @test all(grid.c[3].n.h .≈ [grid.corn.z[1,1,k] for k in 1:grid.c[3].n.s])
end


module ITensorGaussianMPS

using ITensors
using LinearAlgebra

import LinearAlgebra: Givens

export slater_determinant_to_mps,
       slater_determinant_to_gmps,
       hopping_hamiltonian,
       slater_determinant_matrix

include("gmps.jl")

end


"""

#### References

Extreme learning machine: Theory and applications.

Guang-Bin Huang, Qin-Yu Zhu, and Chee-Kheong Siew.

Neurocomputing, 2006 vol. 70 (1-3) pp. 489-501.

http://linkinghub.elsevier.com/retrieve/pii/S0925231206000385
"""
mutable struct ELM{TN,TV} <: AbstractSLFN where TN <:AbstractNodeInput where TV <:AbstractArray{Float64}
    p::Int  # Number of training points
    q::Int  # Dimensionality of function domain
    s::Int  # number of neurons
    neuron_type::TN
    μx::Vector{Float64}
    σx::Vector{Float64}
    Wt::Matrix{Float64}  # transpose of W matrix
    d::Vector{Float64}
    v::TV

    function ELM(p::Int, q::Int, s::Int, neuron_type::TN, μx, σx) where TN
        Wt = 2*rand(q, s) .- 1
        d = rand(s)
        new{TN,typeof(d)}(p, q, s, neuron_type, μx, σx, Wt, d)
    end
end

function ELM(x::AbstractArray, y::TV;
             neuron_type::TN=Linear(Tanh()), s::Int=size(x, 1),
             reg::AbstractLinReg=LSSVD()) where TN<:AbstractNodeInput where TV<:AbstractArray
    q = size(x, 2)  # dimensionality of function domain
    p = size(x, 1)  # number of training points
    s = min(p, s)   # Can't have more neurons than training points
    xn, μx, σx = standardize(x[:, :])
    # out = ELM{TN,TV}(p, q, s, neuron_type, μx, σx)
    out = ELM(p, q, s, neuron_type, μx, σx)
    fit!(out, xn, y, reg)
end

## API methods
function fit!(elm::ELM, x::AbstractArray, y::AbstractArray,
              reg::AbstractLinReg=LSSVD())
    S = hidden_out(elm, x)
    elm.v = regress(reg, S, y)
    elm
end

function Base.show(io::IO, elm::ELM{TA}) where TA
    s =
    """
    ELM with
      - $(TA) Activation function
      - $(elm.q) input dimension(s)
      - $(elm.s) neuron(s)
      - $(elm.p) training point(s)
    """
    print(io, s)
end


#TODO make sure to renew V and S at each epoch

function layerUpdateParams!(
    model::Model,
    cLayer::CL,
    cnt::Integer = -1;
    tMiniBatch::Integer = 1,
    kwargs...,
) where {CL<:ConvLayer}

    optimizer = model.optimizer
    α = model.α
    β1, β2, ϵAdam = model.β1, model.β2, model.ϵAdam

    if cLayer.updateCount >= cnt
        return nothing
    end #if cLayer.updateCount >= cnt

    if !all(
        i -> (i.updateCount == cLayer.nextLayers[1].updateCount),
        cLayer.nextLayers,
    )
        return nothing
    end

    #initialize the needed variables to hold the corrected values
    #it is being init here cause these are not needed elsewhere

    if optimizer == :adam || optimizer == :momentum
        VCorrected = Dict(:dw => similar(cLayer.dW), :db => similar(cLayer.dB))

        if tMiniBatch == 1
            cLayer.V[:dw] .= 0
            cLayer.V[:db] .= 0
        end

        cLayer.V[:dw] .= β1 .* cLayer.V[:dw] .+ (1 - β1) .* cLayer.dW
        cLayer.V[:db] .= β1 .* cLayer.V[:db] .+ (1 - β1) .* cLayer.dB

        ##correcting
        VCorrected[:dw] .= cLayer.V[:dw] ./ (1 - β1^tMiniBatch)
        VCorrected[:db] .= cLayer.V[:db] ./ (1 - β1^tMiniBatch)

        if optimizer == :adam
            SCorrected =
                Dict(:dw => similar(cLayer.dW), :db => similar(cLayer.dB))

            if tMiniBatch == 1
                cLayer.S[:dw] .= 0
                cLayer.S[:db] .= 0
            end
            cLayer.S[:dw] .= β2 .* cLayer.S[:dw] .+ (1 - β2) .* (cLayer.dW .^ 2)
            cLayer.S[:db] .= β2 .* cLayer.S[:db] .+ (1 - β2) .* (cLayer.dB .^ 2)

            ##correcting
            SCorrected[:dw] .= cLayer.S[:dw] ./ (1 - β2^tMiniBatch)
            SCorrected[:db] .= cLayer.S[:db] ./ (1 - β2^tMiniBatch)

            ##update parameters with adam
            cLayer.W .-=
                (α .* (VCorrected[:dw] ./ (sqrt.(SCorrected[:dw]) .+ ϵAdam)))
            cLayer.B .-=
                (α .* (VCorrected[:db] ./ (sqrt.(SCorrected[:db]) .+ ϵAdam)))

            VCorrected = nothing
            SCorrected = nothing
        else#if optimizer==:momentum

            cLayer.W .-= (α .* VCorrected[:dw])
            cLayer.B .-= (α .* VCorrected[:db])
            VCorrected = nothing
        end #if optimizer==:adam
    else
        cLayer.W .-= (α .* cLayer.dW)
        cLayer.B .-= (α .* cLayer.dB)
    end #if optimizer==:adam || optimizer==:momentum

    cLayer.updateCount += 1

    # Base.GC.gc()

    return nothing
end #layerUpdateParams!

function layerUpdateParams!(
    model::Model,
    cLayer::PL,
    cnt::Integer = -1;
    tMiniBatch::Integer = 1,
    kwargs...,
) where {PL<:PoolLayer}

    if cLayer.updateCount >= cnt
        return nothing
    end #if cLayer.updateCount >= cnt

    if !all(
        i -> (i.updateCount == cLayer.nextLayers[1].updateCount),
        cLayer.nextLayers,
    )
        return nothing
    end

    cLayer.updateCount += 1

    return nothing
end #updateParams!


using MPIDistributedArrays
using MPI
MPIDistributedArrays.initialize()
comm = MPI.COMM_WORLD
myrank = MPI.Comm_rank(comm)
nprocs = MPIDistributedArrays.Comm_size()

function partial_tr(A::Array{T,3},NC,NV) where {T <: Number}
    s = zero(T)
    for i=1:NV
        for k=1:NC
            s += A[k,k,i]
        end
    end
    return s
end

function loop_tr(A::MArray{T,N,B},NC,NV) where {T <: Number,N,B}
    #i_local = A.localindices
    s = zero(T)
    A2 = A.localpart
    #_,_,N3 = size(A2)
    #NC2 = NC + 2
    #NV2= A.localindices[3][end] - A.localindices[3][1] + 1# length(A.localindices[3])
    for i = 1:NV#length(A.localindices[3])
        for k=1:NC
            s += A2[k,k,i]
        end
    end
    return s
end

function partial_tr(A::MArray{T,N,B},NC,NV) where {T <: Number,N,B}
    A2 = A.localpart
    N3 = A.localsize[3]
    s = partial_tr(A2,NC,N3)
    return MPI.Allreduce(s,MPI.SUM,A.comm)
    #return reduce_sum(A,s)
end

function test()
    NC = 3
    NV = 256*1000

    NVs = 256*[1,10,100,200,500,1000,5000]
    sc = []
    pc = []
    for NV in NVs
        A_normal = rand(ComplexF64,NC,NC,NV)
        #num_workers = nworkers()

        n = 100
        ns = 10
        ks = []
        for k=1:n
            stime = time_ns()*10^(-3)
            tr_normal = partial_tr(A_normal,NC,NV)
            etime = time_ns()*10^(-3)
            push!(ks,(etime-stime))
            if myrank == 0
                #println("normal: $tr_normal")
                #println("$k-th ",(etime-stime)*10^6," micro sec")
            end
        end

        tr_normal = partial_tr(A_normal,NC,NV)
        if myrank == 0
            println(sum(ks[ns:end])/length(ks[ns:end])," micro sec")
            println("NV = $NV, normal: $tr_normal")
        end
        push!(sc,sum(ks[ns:end])/length(ks[ns:end]))

        MPI.Barrier(comm)
        if myrank == 0
            println("MArray")
        end

        A_dist = distribute(A_normal,[1,1,nprocs])

        ks = []

        for k=1:n
            stime = time_ns()*10^(-3)
            #tr_dist = partial_tr(A_dist,NC,NV)
            tr_dist = partial_tr(A_dist,NC,NV)
            #tr_dist = partial_tr(A_normal,NC,NV)
            #tr_normal = partial_tr(A_normal,NC,NV)
            etime =time_ns()*10^(-3)
            push!(ks,(etime-stime))
            if myrank == 0
                #println("MArray: $tr_dist")
                #println("$k-th ",(etime-stime)*10^6," micro msec")
            end
        end

        tr_dist = partial_tr(A_dist,NC,NV)
        if myrank == 0
            println(sum(ks[ns:end])/length(ks[ns:end])," micro sec")
            println("NV = $NV, MArray: $tr_dist")
            #println("MArray: $tr_dist")
        end

        push!(pc,sum(ks[ns:end])/length(ks[ns:end]))
        free(A_dist )
    end

    if myrank == 0
        println(NVs)
        println(sc)
        println(pc)
    end

end
test()


using Parameters, Combinatorics, Distributions

"""
The purpose of this module is to define a flexible stochastic blockmodel for hypergraphs. At this stage, all we are aiming to do is *sample* from the model give a partition, a group intensity function, and a degree vector.
"""

Parameters.@with_kw mutable struct hypergraph
    """
    A very simple hypergraph composite type, designed to hold a node list N, an edge list E, a degree sequence D,
    """

    N::Vector{Int64}
    E::Dict{Int64, Dict}
    D::Array{Int64, 1} = Array{Int64, 1}()

end

function sampleEdge(S::Vector{Int64}, Z::Vector{Int64}, ϑ::Vector{Float64}, Ω::IntensityFunction; α::Any)
    """
    Sample a Poisson number of edges on a sequence of nodes S according to the hypergraph SBM law.
    # Arguments
    S::Vector{Int64},  an array of node indices
    Z::Vector{Int64}, an integer array of length n. Each integer is a group label.
    Ω::IntensityFunction, an IntensityFunction governing interaction rates between nodes of different groups.
    ϑ::Vector{Float64},  a nonnegative float array of length n
    # Returns
    x::Int64, an integer number of edge counts sampled according to a Poisson random variable with specified formula.
    """
    z = Z[S]
    θ = ϑ[S]

    # combinatorial factor associated with repeated indices. Equivalent to sampling a separate Poisson for each of the c distinct permutations of the node labels S
    c = counting_coefficient(S)
    X = Distributions.Poisson(prod(θ)*Ω.ω(Ω.P(z), α)*c)
    return(rand(X))
end

function sampleEdges(Z::Vector{Int64}, ϑ::Vector{Float64}, Ω::IntensityFunction; α::Vector{Float64}, kmax::Int64=3, kmin::Int64=2)
    """
    Sample edges of specified sizes according to the degree-corrected hypergraph configuration model (DCHSBM). Runs sampleEdge() for each possible distinct sequence of nodes of size between kmin and kmax, inclusive.
    # Arguments
    Z::Vector{Int64}, a vector of group labels
    ϑ::Vector{Float64}, a vector of degree parameters
    Ω::IntensityFunction, an intensity function governing the rate of interaction between nodes of various groups.
    α::Vector{Float64}, a parameter passed to Ω
    kmax::Int64, the size of the largest hyperedge to form
    kmin::Int64, the size of the smallest hyperedge to form
    # Returns
    E::Dict{Int64,Dict{Vector{Int64}, Int64}}, a dictionary keyed by edge sizes. E[k] is itself a dictionary whose keys are ordered sets of nodes and whose values are the number of edges on those node sets.
    N::Vector{Int64}, the list of nodes.
    """
    E = Dict{Int64, Dict}() # initialize empty edge list
    n = length(Z)           # number of nodes inferred from group labels

    for k in kmin:kmax
        T = Combinatorics.with_replacement_combinations(1:n, k) # all distinct combos of k nodes
        Ek = Dict{Vector{Int64}, Int64}()         # list of edges of size k
        for S in T
            X = sampleEdge(S, Z, ϑ, Ω; α=α)
            if X > 0                              # only store combos with at least one edge
                Ek[S] = X
            end
        end
        E[k] = Ek
    end
    N = 1:n
    return(E, N)
end

function sampleEdges(Z::Dict, ϑ::Dict, Ω::IntensityFunction; α::Any, kmax::Int64=3, kmin::Int64=2)
    """
    An alternative method when Z and ϑ are specified as dicts keyed by node rather than as arrays.
    """
    Z = [Z[i] for i in 1:length(Z)]
    ϑ = [ϑ[i] for i in 1:length(ϑ)]
    sampleEdges(Z, ϑ, Ω; α=α, kmax=kmax, kmin=kmin)
end

function computeDegrees(E::Dict{Int64, Dict}, N::Vector{Int64})
    """
    Compute the degree sequence of an edge list.
    """

    d = zeros(length(N))

    for k in keys(E)
        Ek = E[k]
        for e in keys(Ek)
            for i in e
                d[i] += 1
            end
        end
    end
    return(d)
end

function computeDegrees(H::hypergraph)
    return computeDegrees(H.E, H.N)
end

function computeDegrees!(H::hypergraph)
    """
    Compute the degree sequence of a hypergraph and store it as a field of the hypergraph.
    """
    H.D = computeDegrees(H)
end

function sampleSBM(args...;kwargs...)
    """
    Sample a hypergraph with specified parameters and return it with its degree sequence pre-computed. This is the primary user-facing function for sampling tasks.
    # Arguments
    args and kwargs: See documentation for sampleEdges()
    # Returns
    H::hypergraph, a hypergraph sampled according to the DCHSBM with specified args and kwargs.
    """
    E, N = sampleEdges(args...;kwargs...)
    H = hypergraph(E = E, N = N)

    computeDegrees!(H)
    return(H)
end

function countEdges(H::hypergraph)
    """
    count the number of edges in H
    """
    sum([length(H.E[k]) for k in keys(H.E)])
end

Base.copy(H::hypergraph) = hypergraph(H.N, H.E, H.D)