text:
#Testbed for synthetic spike recovery from example background
using BioBackgroundModels, BioMotifInference, Random, Distributed, Distributions, Serialization
Random.seed!(786)
#CONSTANTS
no_obs=500
obsl=100:200

const folder_path="/bench/PhD/NGS_binaries/BBM/refined_folders"

report_folders=deserialize(folder_path)

bhmm_vec=Vector{BHMM}()
push!(bhmm_vec,report_folders["intergenic"].partition_report.best_model[2])
push!(bhmm_vec,report_folders["periexonic"].partition_report.best_model[2])
push!(bhmm_vec,report_folders["exon"].partition_report.best_model[2])

bhmm_dist=Categorical([.6,.2,.2])

struc_sig_1=[.1 .7 .1 .1
           .1 .1 .1 .7
           .1 .7 .1 .1]
struc_sig_2=[.1 .1 .7 .1
           .1 .1 .7 .1
           .7 .1 .1 .1]
periodicity=8
struc_frac_obs=.75

tata_box=[.05 .05 .05 .85
          .85 .05 .05 .05
          .05 .05 .05 .85
          .85 .05 .05 .05
          .425 .075 .075 .425
          .85 .05 .05 .05
          .425 .075 .075 .425]
caat_box=[.15 .15 .55 .15
          .15 .15 .55 .15
          .05 .85 .05 .05
          .05 .85 .05 .05
          .85 .05 .05 .05
          .85 .05 .05 .05
          .05 .05 .05 .85
          .15 .55 .15 .15
          .15 .15 .15 .55]
motif_frac_obs=.7
motif_recur_range=1:4

@info "Constructing synthetic sample set 1..."
obs1, bg_scores1, hmm_truth1, spike_truth1 = synthetic_sample(no_obs,obsl,bhmm_vec,bhmm_dist,[struc_sig_1,tata_box],[(true,(struc_frac_obs,periodicity)),(false,(motif_frac_obs,motif_recur_range))])

@info "Constructing synthetic sample set 2..."
obs2, bg_scores2, hmm_truth2, spike_truth2 = synthetic_sample(no_obs,obsl,bhmm_vec,bhmm_dist,[struc_sig_2,caat_box],[(true,(struc_frac_obs,periodicity)),(false,(motif_frac_obs,motif_recur_range))])

@info "Assembling combined sample set..."
obs3=hcat(obs1,obs2)
bg_scores3=hcat(bg_scores1, bg_scores2)

@info "Assembling worker pool..."

#DISTRIBUTED CLUSTERs CONSTANTS
remote_machine = "10.0.0.3"
no_local_processes = 2
no_remote_processes = 6

@info "Spawning local cluster workers..."
worker_pool=addprocs(no_local_processes, topology=:master_worker)
remote_pool=addprocs([(remote_machine,no_remote_processes)], tunnel=true, topology=:master_worker)

worker_pool=vcat(worker_pool, remote_pool)

@info "Loading worker libraries everywhere..."
@everywhere using BioMotifInference, Random
@everywhere Random.seed!(myid())

e1 = "/bench/PhD/NGS_binaries/BMI/e1"
e2 = "/bench/PhD/NGS_binaries/BMI/e2"
e3 = "/bench/PhD/NGS_binaries/BMI/e3"

#JOB CONSTANTS
const ensemble_size = 250
const no_sources = 2
const source_min_bases = 3
const source_max_bases = 12
const source_length_range= source_min_bases:source_max_bases
const mixing_prior = .5
const models_to_permute = ensemble_size * 3
funcvec=full_perm_funcvec
push!(funcvec, BioMotifInference.permute_source)
push!(funcvec, BioMotifInference.permute_source)
args=[Vector{Tuple{Symbol,Any}}() for i in 1:length(funcvec)]
args[end-1]=[(:weight_shift_freq,0.),(:length_change_freq,1.),(:length_perm_range,1:1)]
args[end]=[(:weight_shift_freq,.1),(:length_change_freq,0.),(:weight_shift_dist,Uniform(.00001,.01))]

instruct = Permute_Instruct(funcvec, ones(length(funcvec))./length(funcvec),models_to_permute,100, .02; args=args)

@info "Assembling source priors..."
prior_array= Vector{Matrix{Float64}}()
source_priors = BioMotifInference.assemble_source_priors(no_sources, prior_array)

@info "Assembling ensemble 1..."
isfile(e1*"/ens") ? (ens1 = deserialize(e1*"/ens")) :
    (ens1 = IPM_Ensemble(worker_pool, e1, ensemble_size, source_priors, (falses(0,0), mixing_prior), bg_scores1, obs1, source_length_range))

@info "Converging ensemble 1..."
logZ1 = BioMotifInference.converge_ensemble!(ens1, instruct, worker_pool, .001, backup=(true,250), wk_disp=false, tuning_disp=true, ens_disp=false, conv_plot=true, src_disp=true, lh_disp=false, liwi_disp=false)

@info "Assembling ensemble 2..."
isfile(e2*"/ens") ? (ens1 = deserialize(e2*"/ens")) :
    (ens2 = IPM_Ensemble(worker_pool, e2, ensemble_size, source_priors, (falses(0,0), mixing_prior), bg_scores2, obs2, source_length_range))

@info "Converging ensemble 2..."
logZ2 = BioMotifInference.converge_ensemble!(ens2, instruct, worker_pool, .001, backup=(true,250), wk_disp=false, tuning_disp=true, ens_disp=false, conv_plot=true, src_disp=true, lh_disp=false, liwi_disp=false)

@info "Assembling ensemble 3..."
isfile(e3*"/ens") ? (ens1 = deserialize(e3*"/ens")) :
    (ens3 = IPM_Ensemble(worker_pool, e3, ensemble_size, source_priors, (falses(0,0), mixing_prior), bg_scores3, obs3, source_length_range))

@info "Converging ensemble 3..."

logZ3 = BioMotifInference.converge_ensemble!(ens3, instruct, worker_pool, .001, backup=(true,250), wk_disp=false, tuning_disp=true, ens_disp=false, conv_plot=true, src_disp=true, lh_disp=false, liwi_disp=false)

@info "Evidence ratio of joint 1&2 / 3:"
ratio=exp((logZ1+logZ2)-logZ3)
println(ratio)

rmprocs(worker_pool)


text:
#------------------------------------
# Common Types
#------------------------------------
const TaskIdType = UInt64

const NoTask = TaskIdType(0)
taskid(id::TaskIdType) = id
taskid(th::Thunk) = TaskIdType(th.id)
#taskid(ch::Chunk) = TaskIdType(hash(ch))
#taskid(executable) = TaskIdType(hash(executable))

mutable struct NodeEnv
    brokerid::UInt64
    host::IPAddr
    broker_task::Union{Future,Void}
    executorids::Vector{UInt64}
    executor_tasks::Vector{Future}
    last_task_stat::Vector{Tuple{UInt64,Future}}

    function NodeEnv(brokerid::Integer, host::IPAddr, executorids::Vector{Int})
        new(brokerid, host, nothing, executorids, Vector{Future}(), Vector{Tuple{UInt64,Future}}())
    end
end

# NodeEnv list, optionally accompanied by the nodehash
const NodeEnvList = Union{Vector{NodeEnv},Tuple{UInt64,Vector{NodeEnv}}}

mutable struct RunEnv
    rootpath::String                # metadata path to work with
    masterid::UInt64
    nodes::Vector{NodeEnv}          # process tree with designated brokers and executors
    nodehash::UInt64                # used to detect changes in process tree
    reset_task::Union{Task,Void}    # async reset task fired after each run
    debug::Bool
    profile::Bool
    remotetrack::Bool

    function RunEnv(; rootpath::String="/D", masterid::Int=myid(), nodes::NodeEnvList=setup_nodes(masterid), debug::Bool=false, profile::Bool=false, remotetrack::Bool=false)
        env = new(rootpath, masterid, Vector{NodeEnv}(), 0, nothing, debug, profile, remotetrack)
        reset(env, nodes)
        env
    end
end

function setupfdbrc(refpath::Vector{UInt8})
    DagScheduler.FdbRC.set_fdb_db()
    empty!(DagScheduler.FdbRC.RefCountRoot)
    append!(DagScheduler.FdbRC.RefCountRoot, refpath)
    nothing
end

function reset(env::RunEnv, nodes::Vector{Int})
    (hash(nodes) === env.nodehash) || reset(env)
    nothing
end

function reset(env::RunEnv, nodes::NodeEnvList=setup_nodes(env.masterid))
    if isa(nodes, Tuple)
        env.nodehash, env.nodes = nodes
    else
        env.nodes = nodes
        env.nodehash = hash(procs())
    end
    nexecutors = 0
    for node in env.nodes
        nw = length(node.executorids)
        (nw > 1) || error("need at least two workers on each node")
        nexecutors += nw
    end
    (nexecutors > 1) || error("need at least two workers")

    # ensure clean paths
    delete_meta(UInt64(env.masterid), env.nodes, env.rootpath)

    @everywhere begin
        MemPool.enable_who_has_read[] = false
        empty!(MemPool.wrkrips)
        Dagger.use_shared_array[] = false
    end

    refpath = convert(Vector{UInt8}, joinpath(env.rootpath, "ref"))
    @sync for w in procs()
        @async remotecall_wait(setupfdbrc, w, refpath)
    end

    nothing
end

function show(io::IO, env::RunEnv)
    attribs = []
    env.debug && push!(attribs, "debug")
    env.profile && push!(attribs, "profile")
    env.remotetrack && push!(attribs, "remotetrack")
    print(io, "RunEnv", isempty(attribs) ? "" : (" (" * join(attribs, ", ") * ")"))
    for node in env.nodes
        print(io, "\n ↳ ", node)
    end
    nothing
end

function show(io::IO, node::NodeEnv)
    print(io, node.host, " [", node.brokerid, " ⇄ ", join(map(string, node.executorids), ","), "]")
    nothing
end

struct ComputeCost
    cpu::Float64        # cost of computing
    net::Float64        # cost of transfering data
end

fval(c::ComputeCost) = (c.cpu + c.net)

function +(c1::ComputeCost, c2::ComputeCost)
    ComputeCost(c1.cpu + c2.cpu, c1.net + c2.net)
end
+(c1::ComputeCost) = c1
isless(c1::ComputeCost, c2::ComputeCost) = fval(c1) < fval(c2)
/(c1::ComputeCost, c2::ComputeCost) = fval(c1) / fval(c2)

# here is where the network cost can be adjusted to account for the computation that this Thunk does
# if it is a reduction step, apply the reduction before summing the input costs
cost_function(t::Thunk) = sum

# consider splitting node for scheduling only if there is no clear association with any node
const SPLIT_COST_RATIO = 1.5
# consider splitting node for scheduling only if the compute cost is substantial
const SPLIT_CAPACITY_RATIO = 0.5

mean_node_capacity(runenv::RunEnv) = mean([length(node.executorids)*1 for node in runenv.nodes])

#------------------------------------
# Logging utilities
#------------------------------------
function tasklog(env, msg...)
    env.debug && info(now(), " ", env.role, " : ", env.id, " : ", env.brokerid, " : ", msg...)
    env.remotetrack && logmsg(Symbol("t"*string(object_id(current_task()))), now(), " ", env.role, " : ", env.id, " : ", env.brokerid, " : ", msg...)
    nothing
end

function taskexception(env, ex, bt)
    xret = CapturedException(ex, bt)
    tasklog(env, "exception ", xret)
    @show xret
    xret
end

function profile_init(prof::Bool, myid::String)
    if prof
        Profile.init(10^7, 0.01)
        Profile.start_timer()
    end
    nothing
end

function profile_end(prof::Bool, myid::String)
    if prof
        Profile.stop_timer()
        open("/tmp/$(myid).profile", "w") do f
            Profile.print(IOContext(f, :displaysize => (256, 1024)))
        end
        Profile.clear()
    end
    nothing
end

function remotetrack_init(remotetrack::Bool)
    remotetrack && start_sender()
    nothing
end

function remotetrack_end(remotetrack::Bool)
    remotetrack && stop_sender()
    nothing
end

#------------------------------------
# Topology utilities
#------------------------------------
function worker_distribution()
    d = Dict{IPv4,Vector{Int}}()
    for w in Base.Distributed.PGRP.workers
        wip = IPv4(isa(w, Base.Distributed.Worker) ? get(w.config.bind_addr) : w.bind_addr)
        if wip in keys(d)
            push!(d[wip], w.id)
        else
            d[wip] = [w.id]
        end
    end

    # disambiguate localhost
    loc = ip"127.0.0.1"
    if loc in keys(d)
        realip = remotecall_fetch(getipaddr, first(d[loc]))
        if realip in keys(d)
            append!(d[realip], d[loc])
        else
            d[realip] = d[loc]
        end
        delete!(d, loc)
    end
    d
end

function setup_nodes(master=1)
    wd = worker_distribution()
    nodehash = hash(sort(vcat(collect(values(wd))...)))
    frefservers = Dict{IPv4,Vector{Int}}()

    if length(wd) == 1
        # only one node. make master the broker and use shm meta_store
        @everywhere begin
            eval(:(using DagScheduler))
            eval(:(using Dagger))
            eval(:(using MemPool))
            DagScheduler.META_IMPL[:node] = "DagScheduler.ShmemMeta.ShmemExecutorMeta"
            DagScheduler.META_IMPL[:cluster] = "DagScheduler.ShmemMeta.ShmemExecutorMeta"
        end
        executors = first(values(wd))
        splice!(executors, findfirst(executors, master))
        nodes = [NodeEnv(master, first(keys(wd)), executors)]
        frefservers[first(keys(wd))] = [master]
    else
        # multiple nodes. select one broker on each node (minimum pid on each group is the broker)
        @everywhere begin
            eval(:(using DagScheduler))
            eval(:(using Dagger))
            eval(:(using MemPool))
            DagScheduler.META_IMPL[:node] = "DagScheduler.ShmemMeta.ShmemExecutorMeta"
            DagScheduler.META_IMPL[:cluster] = "DagScheduler.SimpleMeta.SimpleExecutorMeta"
        end
        nodes = NodeEnv[]
        for (ip,wrkrs) in wd
            sort!(filter!(w->w!=master, wrkrs))
            if !isempty(wrkrs)
                brokerid = shift!(wrkrs)
                # include a node only if it has at least 2 executors
                (length(wrkrs) > 1) && push!(nodes, NodeEnv(brokerid, ip, wrkrs))
                frefservers[ip] = [brokerid]
            else
                frefservers[ip] = [master]
            end
        end
    end

    @sync for w in procs()
        @async remotecall_wait(w) do
            MemPool.enable_random_fref_serve[] = false
            empty!(MemPool.wrkrips)
            merge!(MemPool.wrkrips, frefservers)
        end
    end

    nodehash, nodes
end

#------------------------------------
# DAG Utilities
#------------------------------------
function collect_chunks(dag)
    if isa(dag, Chunk)
        return collect(dag)
    elseif isa(dag, Thunk)
       dag.inputs = map(collect_chunks, dag.inputs)
       return dag
    else
       return dag
    end
end

function get_drefs(dag, bucket::Vector{Chunk}=Vector{Chunk}())
   if isa(dag, Chunk)
       if isa(dag.handle, DRef)
           push!(bucket, dag)
       end
   elseif isa(dag, Thunk)
       map(x->get_drefs(x, bucket), dag.inputs)
   end
   bucket
end

#=
get_frefs(dag) = map(chunktodisk, get_drefs(dag))
=#

chunktodisk(chunks::Vector{Chunk}) = map(chunktodisk, chunks)
chunktodisk(chunk::Chunk) = isa(chunk.handle, DRef) ? Chunk{Any,FileRef}(chunk.chunktype, chunk.domain, movetodisk(chunk.handle), chunk.persist) : chunk

function walk_dag(fn, dag_node, update::Bool, stop_nodes=TaskIdType[], depth::Int=1)
    if isa(dag_node, Thunk) && !(taskid(dag_node) in stop_nodes)
        if update
            dag_node.inputs = map(x->walk_dag(fn, x, update, stop_nodes, depth+1), dag_node.inputs)
        else
            for inp in dag_node.inputs
                walk_dag(fn, inp, update, stop_nodes, depth+1)
            end
        end
    end
    fn(dag_node, depth)
end

function filter!(fn, dag_node::Thunk)
    for inp in dag_node.inputs
        isa(inp, Thunk) && filter!(fn, inp)
    end
    newinputs = filter!(fn, collect(dag_node.inputs))
    if length(newinputs) < length(dag_node.inputs)
        dag_node.inputs = tuple(newinputs...)
    end
    dag_node
end

function depth(dag_node)
    let maxd = 1
        f = (n,d)->begin
            maxd = max(maxd, d)
        end
        walk_dag(f, dag_node, false)
    end
end

function find_node(dag_node, id::TaskIdType)
    let node = nothing
        f = (n,d)->begin
            if istask(n) && (taskid(n) === id)
                node = n
            end
            node
        end
        walk_dag(f, dag_node, false)
    end
end
find_node(id::TaskIdType) = Dagger._thunk_dict[id]

ischild(dag_node, id::TaskIdType) = (nothing !== find_node(dag_node, id))

persist_chunks!(dag) = walk_dag(dag, true) do node,depth
    if isa(node, Chunk)
        node.persist = true
    end
    node
end

dref_to_fref(dag) = dref_to_fref!(deepcopy(dag))
dref_to_fref!(dag) = walk_dag(dag, true) do node,depth
    if isa(node, Chunk) && isa(node.handle, DRef)
        chunktodisk(node)
    else
        node
    end
end

ext: jl


text:
################################################################################
#
#  Reduction of sparse rows modulo sparse upper triangular matrices
#
################################################################################

@doc Markdown.doc"""
    reduce(A::SMat{T}, g::SRow{T}) -> SRow{T}

Given an upper triangular matrix $A$ over a field and a sparse row $g$, this
function reduces $g$ modulo $A$.
"""
function reduce(A::SMat{T}, g::SRow{T}) where {T <: FieldElement}
  return _reduce_field(A, g)
end

function reduce(A::SMat{nmod}, g::SRow{nmod})
  return _reduce_field(A, g)
end

function _reduce_field(A::SMat{T}, g::SRow{T}) where {T}
  @hassert :HNF 1  isupper_triangular(A)
  #assumes A is upper triangular, reduces g modulo A
  # supposed to be a field...
  if A.r == A.c
    return sparse_row(base_ring(A))
  end
  while length(g)>0
    s = g.pos[1]
    j = 1
    while j<= nrows(A) && A.rows[j].pos[1] < s
      j += 1
    end
    if j > nrows(A) || A.rows[j].pos[1] > s
      break
    end
    @hassert :HNF 2  A.rows[j].pos[1] == g.pos[1]
    @hassert :HNF 2  (A.rows[j].values[1]) == 1
    p = g.values[1]
    g = Hecke.add_scaled_row(A[j], g, -p)
    @hassert :HNF 2  length(g)==0 || g.pos[1] > A[j].pos[1]
  end
  if length(g) > 0
    li = inv(g.values[1])
    for i=1:length(g)
      g.values[i] *= li
    end
  end
  return g
end

@doc Markdown.doc"""
    reduce(A::SMat{fmpz}, g::SRow{fmpz}) -> SRow{fmpz}

Given an upper triangular matrix $A$ over a field and a sparse row $g$, this
function reduces $g$ modulo $A$.
"""
function reduce(A::SMat{fmpz}, g::SRow{fmpz})
  @hassert :HNF 1  isupper_triangular(A)
  #assumes A is upper triangular, reduces g modulo A
  #until the 1st (pivot) change in A
  new_g = false
  while length(g)>0
    s = g.pos[1]
    j = 1
    while j<= nrows(A) && A.rows[j].pos[1] < s
      j += 1
    end
    if j > nrows(A) || A.rows[j].pos[1] > s
      if g.values[1] < 0
        if !new_g
          g = copy(g)
        end
        for i=1:length(g.values)
          g.values[i] *= -1
        end
      end
      return g
    end
    p = g.values[1]
    if divides(p, A.rows[j].values[1])[1]
      g = Hecke.add_scaled_row(A[j], g, - divexact(p, A.rows[j].values[1]))
      new_g = true
      @hassert :HNF 2  length(g)==0 || g.pos[1] > A[j].pos[1]
    else
      x, a, b = gcdx(A.rows[j].values[1], p)
      @hassert :HNF 2  x > 0
      c = -div(p, x)
      d = div(A.rows[j].values[1], x)
      A[j], g = Hecke.transform_row(A[j], g, a, b, c, d)
      new_g = true
      @hassert :HNF 2  A[j].values[1] == x
      @hassert :HNF 2  length(g)==0 || g.pos[1] > A[j].pos[1]
    end
  end

  if length(g.values) > 0 && g.values[1] < 0
    if !new_g
      g = copy(g)
    end
    for i=1:length(g.values)
      g.values[i] *= -1
    end
  end
  return g
end

@doc Markdown.doc"""
    reduce(A::SMat{fmpz}, g::SRow{fmpz}, m::fmpz) -> SRow{fmpz}

Given an upper triangular matrix $A$ over the integers, a sparse row $g$ and an
integer $m$, this function reduces $g$ modulo $A$ and returns $g$
modulo $m$ with respect to the symmetric residue system.
"""
function reduce(A::SMat{fmpz}, g::SRow{fmpz}, m::fmpz)
  @hassert :HNF 1  isupper_triangular(A)
  #assumes A is upper triangular, reduces g modulo A
  g = copy(g)
  mod_sym!(g, m)
  while length(g)>0
    s = g.pos[1]
    j = 1
    while j<= nrows(A) && A.rows[j].pos[1] < s
      j += 1
    end
    if j > nrows(A) || A.rows[j].pos[1] > s
      if mod_sym(g.values[1], m) < 0
        for i=1:length(g.values)
          g.values[i] *= -1
        end
        @hassert :HNF 3  mod_sym(g.values[1], m) > 0
      end
      mod_sym!(g, m)
      return g
    end
    st_g = 2
    st_A = 2
    p = g.values[1]
    if divides(p, A.rows[j].values[1])[1]
      g = add_scaled_row(A[j], g, - divexact(p, A.rows[j].values[1]))
      mod_sym!(g, m)
      @hassert :HNF 2  length(g)==0 || g.pos[1] > A[j].pos[1]
    else
      x, a, b = gcdx(A.rows[j].values[1], p)
      @hassert :HNF 2  x > 0
      c = -div(p, x)
      d = div(A.rows[j].values[1], x)
      A[j], g = Hecke.transform_row(A[j], g, a, b, c, d)
      new_g = true
#      @hassert :HNF 2  A[j].values[1] == x
      mod_sym!(g, m)
      mod_sym!(A[j], m)
#      @hassert :HNF 2  length(g)==0 || g.pos[1] > A[j].pos[1]
#      @hassert :HNF 2  A[j].values[1] > 0
    end
  end
  if length(g.values) > 0 && mod_sym(g.values[1], m) < 0
    for i=1:length(g.values)
      g.values[i] *= -1
    end
  end
  Hecke.mod_sym!(g, m)
#  @hassert :HNF 1  length(g.pos) == 0 || g.values[1] >= 0
  return g
end

################################################################################
#
#  Saturation
#
################################################################################

@doc Markdown.doc"""
    saturate(A::SMat{fmpz}) -> SMat{fmpz}

Computes the saturation of $A$, that is, a basis for $\mathbf{Q}\otimes M \meet
\mathbf{Z}^n$, where $M$ is the row span of $A$ and $n$ the number of rows of
$A$.

Equivalently, return $TA$ for an invertible rational matrix $T$, such that $TA$
is integral and the elementary divisors of $TA$ are all trivial.
"""
function saturate(A::SMat{fmpz})
  Hti = transpose(hnf(transpose(A)))
  Hti = sub(Hti , 1:nrows(Hti), 1:nrows(Hti))
  Hti = transpose(Hti)
  S, s = solve_ut(Hti, transpose(A))
  @assert isone(s)
  SS = transpose(S)
  return SS
end

################################################################################
#
#  Hermite normal form using Kannan-Bachem algorithm
#
################################################################################

@doc Markdown.doc"""
    find_row_starting_with(A::SMat, p::Int) -> Int

Tries to find the index $i$ such that $A_{i,p} \neq 0$ and $A_{i, p-j} = 0$
for all $j > 1$. It is assumed that $A$ is upper triangular.
If such an index does not exist, find the smallest index larger.
"""
function find_row_starting_with(A::SMat, p::Int)
#  @hassert :HNF 1  isupper_triangular(A)
  start = 0
  stop = nrows(A)+1
  while start < stop - 1
    mid = div((stop + start), 2)
    if A[mid].pos[1] == p
      return mid
    elseif A[mid].pos[1] < p
      start = mid
    else
      stop = mid
    end
  end
  return stop
end

# If trafo is set to Val{true}, then additionaly an Array of transformations
# is returned.
function reduce_up(A::SMat{fmpz}, piv::Array{Int, 1},
                                  trafo::Type{Val{N}} = Val{false}) where N

  with_transform = (trafo == Val{true})
  if with_transform
    trafos = SparseTrafoElem{fmpz, fmpz_mat}[]
  end

  sort!(piv)
  p = find_row_starting_with(A, piv[end])

  for red=p-1:-1:1
    # the last argument should be the smallest pivot larger then pos[1]
    if with_transform
      A[red], new_trafos = reduce_right(A, A[red], max(A[red].pos[1]+1, piv[1]), trafo)
      for t in new_trafos
        t.j = red
      end
      append!(trafos, new_trafos)
    else
      A[red] = reduce_right(A, A[red], max(A[red].pos[1]+1, piv[1]))
    end
  end
  with_transform ? (return trafos) : nothing
end

# If trafo is set to Val{true}, then additionaly an Array of transformations
# is returned.
@doc Markdown.doc"""
    reduce_full(A::SMat{fmpz}, g::SRow{fmpz},
                          trafo = Val{false}) -> SRow{fmpz}, Vector{Int}

Reduces $g$ modulo $A$ and assumes that $A$ is upper triangular.

The second return value is the array of pivot elements of $A$ that changed.

If `trafo` is set to `Val{true}`, then additionally an array of transformations
is returned.
"""
function reduce_full(A::SMat{fmpz}, g::SRow{fmpz}, trafo::Type{Val{T}} = Val{false}) where T
#  @hassert :HNF 1  isupper_triangular(A)
  #assumes A is upper triangular, reduces g modulo A

  with_transform = (trafo == Val{true})
  no_trafo = (trafo == Val{false})

  if with_transform
    trafos = SparseTrafoElem{fmpz, fmpz_mat}[]
  end

  new_g = false

  piv = Int[]
  while length(g)>0
    s = g.pos[1]
    j = 1
    while j<= nrows(A) && A.rows[j].pos[1] < s
      j += 1
    end
    if j > nrows(A) || A.rows[j].pos[1] > s
      if g.values[1] < 0
        # Multiply row g by -1
        if with_transform
          push!(trafos, sparse_trafo_scale(nrows(A) + 1, fmpz(-1)))
        end
        if !new_g
          g = copy(g)
        end
        for i=1:length(g.values)
          g.values[i] *= -1
        end
      end

      if with_transform
        g, new_trafos  = reduce_right(A, g, 1, trafo)
        append!(trafos, new_trafos)
      else
        g = reduce_right(A, g)
      end
      new_g = true

      if A.r == A.c
        @hassert :HNF 1  length(g) == 0 || minimum(g) >= 0
      end

      with_transform ? (return g, piv, trafos) : (return g, piv)

    end
    p = g.values[1]
    if divides(p, A.rows[j].values[1])[1]
      sca =  -divexact(p, A.rows[j].values[1])
      g = Hecke.add_scaled_row(A[j], g, sca)
      new_g = true
      with_transform ? push!(trafos, sparse_trafo_add_scaled(j, nrows(A) + 1, sca)) : nothing
      @hassert :HNF 1  length(g)==0 || g.pos[1] > A[j].pos[1]
    else
      x, a, b = gcdx(A.rows[j].values[1], p)
      @hassert :HNF 1  x > 0
      c = -div(p, x)
      d = div(A.rows[j].values[1], x)
      A[j], g = Hecke.transform_row(A[j], g, a, b, c, d)
      new_g = true
      if with_transform
        push!(trafos, sparse_trafo_para_add_scaled(j, nrows(A) + 1, a, b, c, d))
      end
      @hassert :HNF 1  A[j].values[1] == x
      @hassert :HNF 1  length(g)==0 || g.pos[1] > A[j].pos[1]
      push!(piv, A[j].pos[1])
      if with_transform
        A[j], new_trafos = reduce_right(A, A[j], A[j].pos[1]+1, trafo)
        # We are updating the jth row
        # Have to adjust the transformations
        for t in new_trafos
          t.j = j
        end
        # Now append
        append!(trafos, new_trafos)
      else
        A[j] = reduce_right(A, A[j], A[j].pos[1]+1, trafo)
      end

      if A.r == A.c
        @hassert :HNF 1  minimum(A[j]) >= 0
      end
    end
  end
  if length(g.values) > 0 && g.values[1] < 0
    if !new_g
      g = copy(g)
    end
    for i=1:length(g.values)
      g.values[i] *= -1
    end
    with_transform ? push!(trafos, sparse_trafo_scale!{fmpz}(nrows(A) + 1, fmpz(-1))) : nothing
  end
  if with_transform
    g, new_trafos = reduce_right(A, g, 1, trafo)
    append!(trafos, new_trafos)
  else
    g = reduce_right(A, g)
  end
  if A.r == A.c
    @hassert :HNF 1  length(g) == 0 || minimum(g) >= 0
  end
  with_transform ? (return g, piv, trafos) : (return g, piv)
end

function reduce_right(A::SMat{fmpz}, b::SRow{fmpz},
                      start::Int = 1, trafo::Type{Val{N}} = Val{false}) where N
  with_transform = (trafo == Val{true})
  with_transform ? trafos = [] : nothing
  new = true
  if length(b.pos) == 0
    with_transform ? (return b, trafos) : return b
  end
  j = 1
  while j <= length(b.pos) && b.pos[j] < start
    j += 1
  end
  if j > length(b.pos)
    with_transform ? (return b, trafos) : return b
  end
  p = find_row_starting_with(A, b.pos[j])
  if p > nrows(A)
    with_transform ? (return b, trafos) : return b
  end
  @hassert :HNF 1  A[p] != b
  while j <= length(b.pos)
    while p<nrows(A) && A[p].pos[1] < b.pos[j]
      p += 1
    end
    if A[p].pos[1] == b.pos[j]
      q, r = divrem(b.values[j], A[p].values[1])
      if r < 0
        q -= 1
        r += A[p].values[1]
        @hassert :HNF 1  r >= 0
      end
      if q != 0
        if new
          b = Hecke.add_scaled_row(A[p], b, -q)
          new = false
        else
          Hecke.add_scaled_row!(A[p], b, -q)
        end

        with_transform ? push!(trafos, sparse_trafo_add_scaled(p, nrows(A) + 1, -q)) : nothing
        if r == 0
          j -= 1
        else
          @hassert :HNF 1  b.values[j] == r
        end
      end
    end
    j += 1
  end
  with_transform ? (return b, trafos) : return b
end

@doc Markdown.doc"""
    hnf_extend!(A::SMat{fmpz}, b::SMat{fmpz}, offset::Int = 0) -> SMat{fmpz}

Given a matrix $A$ in HNF, extend this to get the HNF of the concatenation
with $b$.
"""
function hnf_extend!(A::SMat{fmpz}, b::SMat{fmpz}, trafo::Type{Val{N}} = Val{false}; truncate::Bool = false, offset::Int = 0) where N
  rA = nrows(A)
  @vprint :HNF 1 "Extending HNF by:\n"
  @vprint :HNF 1 b
  @vprint :HNF 1 "density $(density(A)) $(density(b))"

  with_transform = (trafo == Val{true})
  with_transform ? trafos = SparseTrafoElem{fmpz, fmpz_mat}[] : nothing

  A_start_rows = nrows(A)  # for the offset stuff

  nc = 0
  for i=b
    if with_transform
      q, w, new_trafos = reduce_full(A, i, trafo)
      append!(trafos, new_trafos)
    else
      q, w = reduce_full(A, i)
    end

    if length(q) > 0
      p = find_row_starting_with(A, q.pos[1])
      if p > length(A.rows)
        # Appending row q to A
        # Do not need to track a transformation
        push!(A, q)
      else
        # Inserting row q at position p
        insert!(A.rows, p, q)
        A.r += 1
        A.nnz += length(q)
        A.c = max(A.c, q.pos[end])
        # The transformation is swapping pairwise from nrows(B) to p.
        # This should be the permutation matrix corresponding to
        # (k k-1)(k-1 k-2) ...(p+1 p) where k = nrows(B)
        if with_transform
          for j in nrows(A):-1:(p+1)
            push!(trafos, sparse_trafo_swap(fmpz, j, j - 1))
          end
        end
      end
      push!(w, q.pos[1])
    else
      # Row i was reduced to zero
      with_transform ? push!(trafos, sparse_trafo_move_row(fmpz, nrows(A) + 1, rA + nrows(b))) : nothing
    end
    if length(w) > 0
      if with_transform
        new_trafos = reduce_up(A, w, trafo)
        append!(trafos, new_trafos)
      else
        reduce_up(A, w)
      end
    end
    @v_do :HNF 1 begin
      if nc % 10 == 0
        println("Now at $nc rows of $(nrows(b)), HNF so far $(nrows(A)) rows")
        println("Current density: $(density(A))")
        println("and size of largest entry: $(nbits(maximum(abs, A))) bits $(sum(nbits, A))")
      end
    end
    nc += 1
  end
  if !truncate
    for i in 1:nrows(b)
      push!(A, sparse_row(base_ring(A)))
    end
  end

  if with_transform && offset != 0
    change_indices!(trafos, A_start_rows, offset)
  end

  with_transform ? (return A, trafos) : (return A)
end

function nbits(s::SRow{fmpz})
  return sum(nbits, s.values)
end

@doc Markdown.doc"""
    hnf_kannan_bachem(A::SMat{fmpz}) -> SMat{fmpz}

Compute the Hermite normal form of $A$ using the Kannan-Bachem algorithm.
"""
function hnf_kannan_bachem(A::SMat{fmpz}, trafo::Type{Val{N}} = Val{false}; truncate::Bool = false) where N
  @vprint :HNF 1 "Starting Kannan Bachem HNF on:\n"
  @vprint :HNF 1 A
  @vprint :HNF 1 "with density $(density(A)); truncating $truncate"

  with_transform = (trafo == Val{true})
  with_transform ? trafos = SparseTrafoElem{fmpz, fmpz_mat}[] : nothing

  B = sparse_matrix(FlintZZ)
  B.c = A.c
  nc = 0
  for i=A
    if with_transform
      q, w, new_trafos = reduce_full(B, i, trafo)
      append!(trafos, new_trafos)
    else
      q, w = reduce_full(B, i)
    end

    if length(q) > 0
      p = find_row_starting_with(B, q.pos[1])
      if p > length(B.rows)
        # Appending row q to B
        # Do not need to track a transformation
        push!(B, q)
      else
        # Inserting row q at position p
        insert!(B.rows, p, q)
        B.r += 1
        B.nnz += length(q)
        B.c = max(B.c, q.pos[end])
        # The transformation is swapping pairwise from nrows(B) to p.
        # This should be the permutation matrix corresponding to
        # (k k-1)(k-1 k-2) ...(p+1 p) where k = nrows(B)
        if with_transform
          for j in nrows(B):-1:(p+1)
            push!(trafos, sparse_trafo_swap(fmpz, j, j - 1))
          end
        end
      end
      push!(w, q.pos[1])
    else
      # Row i was reduced to zero
      with_transform ? push!(trafos, sparse_trafo_move_row(fmpz, nrows(B) + 1, nrows(A))) : nothing
    end
    if length(w) > 0
      if with_transform
        new_trafos = reduce_up(B, w, trafo)
        append!(trafos, new_trafos)
      else
        reduce_up(B, w)
      end
    end
    @v_do :HNF 1 begin
      if nc % 10 == 0
        println("Now at $nc rows of $(nrows(A)), HNF so far $(nrows(B)) rows")
        println("Current density: $(density(B))")
        println("and size of largest entry: $(nbits(maximum(abs, B))) bits")
      end
    end
    nc += 1
  end
  if !truncate
    for i in 1:(nrows(A) - nrows(B))
      push!(B, sparse_row(base_ring(A)))
    end
  end
  with_transform ? (return B, trafos) : (return B)
end

@doc Markdown.doc"""
    hnf(A::SMat{fmpz}) -> SMat{fmpz}

Return the upper right Hermite normal form of $A$.
"""
function hnf(A::SMat{fmpz}; truncate::Bool = false)
  return hnf_kannan_bachem(A, truncate = truncate)
end

@doc Markdown.doc"""
    hnf!(A::SMat{fmpz})

Inplace transform of $A$ into upper right Hermite normal form.
"""
function hnf!(A::SMat{fmpz}; truncate::Bool = false)
  B = hnf(A, truncate = truncate)
  A.rows = B.rows
  A.nnz = B.nnz
  A.r = B.r
  A.c = B.c
end

function reduce_right!(A::SMat{fmpz}, b::SRow{fmpz})
  if length(b.pos) == 0
    return b
  end
  j = 1
  p = find_row_starting_with(A, b.pos[j])
  if p > nrows(A)
    return b
  end
  while j <= length(b.pos)
    while p < nrows(A) && A[p].pos[1] < b.pos[j]
      p += 1
    end
    if A[p].pos[1] == b.pos[j]
      q, r = divrem(b.values[j], A[p].values[1])
      if r < 0
        q -= 1
        r += A[p].values[1]
        @hassert :HNF 1 r >= 0
      end
      if q != 0
        Hecke.add_scaled_row!(A[p], b, -q)
        if r == 0
          j -= 1
        else
          @hassert :HNF 1 b.values[j] == r
        end
      end
    end
    j += 1
  end
  return b
end


text:
using RadeonProRender, GeometryBasics, Colors
const RPR = RadeonProRender

function translationmatrix(t::Vec{3,T}) where {T}
    T0, T1 = zero(T), one(T)
    return Mat{4}(T1, T0, T0, T0, T0, T1, T0, T0, T0, T0, T1, T0, t[1], t[2], t[3], T1)
end

isdefined(Main, :context) && RPR.release(context)

context = RPR.Context()

scene = RPR.Scene(context)
set!(context, scene)

matsys = RPR.MaterialSystem(context, 0)

camera = RPR.Camera(context)
lookat!(camera, Vec3f(1.5, 0.75, 1.5), Vec3f(0), Vec3f(0, 0, 1))
RPR.rprCameraSetFocalLength(camera, 45.0)
set!(scene, camera)

env_light = RadeonProRender.EnvironmentLight(context)
image_path = joinpath(@__DIR__, "studio026.exr")
img = RPR.Image(context, image_path)
set!(env_light, img)
setintensityscale!(env_light, 1.2)
push!(scene, env_light)

light = RPR.PointLight(context)
transform!(light, translationmatrix(Vec3f(2, 0, 5)))
f = 1.5
RPR.setradiantpower!(light, 500 / f, 641 / f, 630 / f)
push!(scene, light)

vol_cube = RadeonProRender.Shape(context, Rect3f(Vec3f(0), Vec3f(1)))
push!(scene, vol_cube)

color_lookup = [colorant"red", colorant"yellow", colorant"green"]
convert(Vector{RGB{Float32}}, color_lookup)
density_lookup = [Vec3f.(10:-1:0);]
n = 128
gridvals = Float32[]
indices1 = UInt64[]

for x in 0:(n-1)
    for y in 0:(n-1)
        for z in 0:(n-1)
            v = y / n
            push!(indices1, (x)*(n*n) + (y)*(n)+z)
            push!(gridvals, v)
        end
    end
end
(a, b) = -1, 2
r = range(-3, stop = 3, length = 800)
z = ((x,y) -> y.^4 - x.^4 + a.*y.^2 + b.*x.^2).(r, r')
me = [cos.(2 .* pi .* sqrt.(x.^2 + y.^2)) .* (4 .* z) for x=r, y=r, z=r]
mini, maxi = extrema(me)
grid = RPR.VoxelGrid(context, (me .- mini) ./ (maxi - mini))
volume1 = RPR.HeteroVolume(context)
RPR.set_albedo_grid!(volume1, grid)
RPR.set_albedo_lookup!(volume1, color_lookup)
RPR.set_density_grid!(volume1, grid)
RPR.set_density_lookup!(volume1, density_lookup)

mat = RPR.MaterialNode(matsys, RPR.RPR_MATERIAL_NODE_TRANSPARENT)
set!(mat, RPR.RPR_MATERIAL_INPUT_COLOR, 1.0f0, 1.0f0, 1.0f0, 1.0f0)
RPR.rprSceneAttachHeteroVolume(scene, volume1)
set!(vol_cube, volume1)
set!(vol_cube, mat)

fb_size = (800, 600)
frame_buffer = RPR.FrameBuffer(context, RGBA, fb_size)
frame_bufferSolved = RPR.FrameBuffer(context, RGBA, fb_size)
set!(context, RPR.RPR_AOV_COLOR, frame_buffer)
# to avoid dark transparency, we raise recursion
RPR.rprContextSetParameterByKey1u(context, RPR.RPR_CONTEXT_MAX_RECURSION, 10)
# set_standard_tonemapping!(context)

begin
    clear!(frame_buffer)
    RPR.rprContextSetParameterByKey1u(context, RPR.RPR_CONTEXT_ITERATIONS, 4)
    RPR.render(context)
    RPR.rprContextResolveFrameBuffer(context, frame_buffer, frame_bufferSolved, true)
    RPR.save(frame_bufferSolved, "test.png")
end


text:
using Arithmetic
using Test

# 4
@testset "Evaluate constant" begin
  @test evaluate(Constant(4)) == 4
end

# 1 * 3
@testset "Simple multiplication" begin
  @test evaluate(BinaryOp(Constant(1), MUL, Constant(3))) == 3
end

# 2 + 4 / 4
@testset "Nested operations" begin
  @test evaluate(BinaryOp(Constant(2), ADD, BinaryOp(Constant(4), DIV, Constant(4)))) == 3
end


text:
# 2N low storage methods
function initialize!(integrator,cache::LowStorageRK2NConstantCache)
  integrator.fsalfirst = integrator.f(integrator.uprev, integrator.p, integrator.t) # Pre-start fsal
  integrator.destats.nf += 1
  integrator.kshortsize = 1
  integrator.k = typeof(integrator.k)(undef, integrator.kshortsize)

  # Avoid undefined entries if k is an array of arrays
  integrator.fsallast = zero(integrator.fsalfirst)
  integrator.k[1] = integrator.fsalfirst
end

@muladd function perform_step!(integrator,cache::LowStorageRK2NConstantCache,repeat_step=false)
  @unpack t,dt,u,f,p = integrator
  @unpack A2end,B1,B2end,c2end = cache

  # u1
  tmp = dt*integrator.fsalfirst
  u   = u + B1*tmp

  # other stages
  for i in eachindex(A2end)
    k = f(u, p, t+c2end[i]*dt)
    integrator.destats.nf += 1
    tmp = A2end[i]*tmp + dt*k
    u   = u + B2end[i]*tmp
  end

  integrator.destats.nf += 1
  integrator.k[1] = integrator.fsalfirst
  integrator.fsalfirst = f(u, p, t+dt) # For interpolation, then FSAL'd
  integrator.u = u
end

function initialize!(integrator,cache::LowStorageRK2NCache)
  @unpack k, tmp, williamson_condition = cache
  integrator.kshortsize = 1
  resize!(integrator.k, integrator.kshortsize)
  integrator.k[1] = k
  integrator.fsalfirst = k # used for get_du
  integrator.fsallast = k
  integrator.f(k ,integrator.uprev,integrator.p,integrator.t) # FSAL for interpolation
  integrator.destats.nf += 1
end

@muladd function perform_step!(integrator,cache::LowStorageRK2NCache,repeat_step=false)
  @unpack t,dt,u,f,p = integrator
  @unpack k,tmp,williamson_condition,stage_limiter!,step_limiter!,thread = cache
  @unpack A2end,B1,B2end,c2end = cache.tab

  # u1
  f(k, u, p, t)
  integrator.destats.nf += 1
  @.. thread=thread tmp = dt*k
  @.. thread=thread u   = u + B1*tmp
  # other stages
  for i in eachindex(A2end)
    if williamson_condition
      f(ArrayFuse(tmp, u, (A2end[i], dt, B2end[i])), u, p, t+c2end[i]*dt)
    else
      @.. thread=thread tmp = A2end[i]*tmp
      stage_limiter!(u, integrator, p, t+c2end[i]*dt)
      f(k, u, p, t+c2end[i]*dt)
      @.. thread=thread tmp = tmp + dt * k
      @.. thread=thread u   = u + B2end[i]*tmp
    end
    integrator.destats.nf += 1
  end
  stage_limiter!(u, integrator, p, t+dt)
  step_limiter!(u, integrator, p, t+dt)
end

# 2C low storage methods
function initialize!(integrator,cache::LowStorageRK2CConstantCache)
  integrator.fsalfirst = integrator.f(integrator.uprev, integrator.p, integrator.t) # Pre-start fsal
  integrator.destats.nf += 1
  integrator.kshortsize = 1
  integrator.k = typeof(integrator.k)(undef, integrator.kshortsize)

  # Avoid undefined entries if k is an array of arrays
  integrator.fsallast = zero(integrator.fsalfirst)
  integrator.k[1] = integrator.fsalfirst
end

@muladd function perform_step!(integrator,cache::LowStorageRK2CConstantCache,repeat_step=false)
  @unpack t,dt,u,f,p = integrator
  @unpack A2end,B1,B2end,c2end = cache

  # u1
  k = integrator.fsalfirst = f(u, p, t)
  integrator.k[1] = integrator.fsalfirst
  integrator.destats.nf += 1
  u = u + B1*dt*k

  # other stages
  for i in eachindex(A2end)
    tmp = u + A2end[i]*dt*k
    k   = f(tmp, p, t+c2end[i]*dt)
    integrator.destats.nf += 1
    u   = u + B2end[i]*dt*k
  end

  integrator.u = u
end

function initialize!(integrator,cache::LowStorageRK2CCache)
  @unpack k,fsalfirst = cache
  integrator.fsalfirst = fsalfirst
  integrator.fsallast = k
  integrator.kshortsize = 1
  resize!(integrator.k, integrator.kshortsize)
  integrator.k[1] = integrator.fsalfirst
  integrator.f(integrator.fsalfirst,integrator.uprev,integrator.p,integrator.t) # FSAL for interpolation
  integrator.destats.nf += 1
end

@muladd function perform_step!(integrator,cache::LowStorageRK2CCache,repeat_step=false)
  @unpack t,dt,u,f,p = integrator
  @unpack k,fsalfirst,tmp = cache
  @unpack A2end,B1,B2end,c2end = cache.tab

  # u1
  @.. k = integrator.fsalfirst
  @.. u = u + B1*dt*k

  # other stages
  for i in eachindex(A2end)
    @.. tmp = u + A2end[i]*dt*k
    f(k, tmp, p, t+c2end[i]*dt)
    integrator.destats.nf += 1
    @.. u   = u + B2end[i]*dt*k
  end

  f(k, u, p, t+dt)
  integrator.destats.nf += 1
end

# 3S low storage methods
function initialize!(integrator,cache::LowStorageRK3SConstantCache)
  integrator.fsalfirst = integrator.f(integrator.uprev, integrator.p, integrator.t) # Pre-start fsal
  integrator.destats.nf += 1
  integrator.kshortsize = 1
  integrator.k = typeof(integrator.k)(undef, integrator.kshortsize)

  # Avoid undefined entries if k is an array of arrays
  integrator.fsallast = zero(integrator.fsalfirst)
  integrator.k[1] = integrator.fsalfirst
end

@muladd function perform_step!(integrator,cache::LowStorageRK3SConstantCache,repeat_step=false)
  @unpack t,dt,uprev,u,f,p = integrator
  @unpack γ12end, γ22end, γ32end, δ2end, β1, β2end, c2end = cache

  # u1
  tmp = u
  u   = tmp + β1*dt*integrator.fsalfirst

  # other stages
  for i in eachindex(γ12end)
    k   = f(u, p, t+c2end[i]*dt)
    integrator.destats.nf += 1
    tmp = tmp + δ2end[i]*u
    u   = γ12end[i]*u + γ22end[i]*tmp + γ32end[i]*uprev + β2end[i]*dt*k
  end

  integrator.fsallast = f(u, p, t+dt) # For interpolation, then FSAL'd
  integrator.destats.nf += 1
  integrator.k[1] = integrator.fsalfirst
  integrator.u = u
end

function initialize!(integrator,cache::LowStorageRK3SCache)
  @unpack k,fsalfirst = cache
  integrator.fsalfirst = fsalfirst
  integrator.fsallast = k
  integrator.kshortsize = 1
  resize!(integrator.k, integrator.kshortsize)
  integrator.k[1] = integrator.fsalfirst
  integrator.f(integrator.fsalfirst,integrator.uprev,integrator.p,integrator.t) # FSAL for interpolation
  integrator.destats.nf += 1
end

@muladd function perform_step!(integrator,cache::LowStorageRK3SCache,repeat_step=false)
  @unpack t,dt,uprev,u,f,p = integrator
  @unpack k,fsalfirst,tmp = cache
  @unpack γ12end, γ22end, γ32end, δ2end, β1, β2end, c2end = cache.tab

  # u1
  @.. tmp = u
  @.. u   = tmp + β1*dt*integrator.fsalfirst

  # other stages
  for i in eachindex(γ12end)
    f(k, u, p, t+c2end[i]*dt)
    integrator.destats.nf += 1
    @.. tmp = tmp + δ2end[i]*u
    @.. u   = γ12end[i]*u + γ22end[i]*tmp + γ32end[i]*uprev + β2end[i]*dt*k
  end

  f(k, u, p, t+dt)
  integrator.destats.nf += 1
end

# 3S+ low storage methods: 3S methods adding another memory location for the embedded method (non-FSAL version)
function initialize!(integrator,cache::LowStorageRK3SpConstantCache)
  integrator.fsalfirst = integrator.f(integrator.uprev, integrator.p, integrator.t) # Pre-start fsal
  integrator.destats.nf += 1
  integrator.kshortsize = 1
  integrator.k = typeof(integrator.k)(undef, integrator.kshortsize)

  # Avoid undefined entries if k is an array of arrays
  integrator.fsallast = zero(integrator.fsalfirst)
  integrator.k[1] = integrator.fsalfirst
end

@muladd function perform_step!(integrator,cache::LowStorageRK3SpConstantCache,repeat_step=false)
  @unpack t,dt,uprev,u,f,p = integrator
  @unpack γ12end, γ22end, γ32end, δ2end, β1, β2end, c2end, bhat1, bhat2end = cache

  # u1
  integrator.fsalfirst = f(uprev, p, t)
  integrator.destats.nf += 1
  integrator.k[1] = integrator.fsalfirst
  tmp = uprev
  u   = tmp + β1*dt*integrator.fsalfirst
  if integrator.opts.adaptive
    utilde = bhat1*dt*integrator.fsalfirst
  end

  # other stages
  for i in eachindex(γ12end)
    k   = f(u, p, t+c2end[i]*dt)
    integrator.destats.nf += 1
    tmp = tmp + δ2end[i]*u
    u   = γ12end[i]*u + γ22end[i]*tmp + γ32end[i]*uprev + β2end[i]*dt*k
    if integrator.opts.adaptive
      utilde = utilde + bhat2end[i]*dt*k
    end
  end

  if integrator.opts.adaptive
    atmp = calculate_residuals(utilde, uprev, u, integrator.opts.abstol, integrator.opts.reltol,integrator.opts.internalnorm,t)
    integrator.EEst = integrator.opts.internalnorm(atmp,t)
  end

  integrator.u = u
end

function initialize!(integrator,cache::LowStorageRK3SpCache)
  @unpack k,fsalfirst = cache
  integrator.fsalfirst = fsalfirst
  integrator.fsallast = k
  integrator.kshortsize = 1
  resize!(integrator.k, integrator.kshortsize)
  integrator.k[1] = integrator.fsalfirst
end

@muladd function perform_step!(integrator,cache::LowStorageRK3SpCache,repeat_step=false)
  @unpack t,dt,uprev,u,f,p = integrator
  @unpack k,tmp,utilde,atmp,stage_limiter!,step_limiter!,thread = cache
  @unpack γ12end, γ22end, γ32end, δ2end, β1, β2end, c2end, bhat1, bhat2end = cache.tab

  # u1
  f(integrator.fsalfirst, uprev, p, t)
  integrator.destats.nf += 1
  @.. thread=thread tmp = uprev
  @.. thread=thread u   = tmp + β1*dt*integrator.fsalfirst
  if integrator.opts.adaptive
    @.. thread=thread utilde = bhat1*dt*integrator.fsalfirst
  end

  # other stages
  for i in eachindex(γ12end)
    stage_limiter!(u, integrator, p, t+c2end[i]*dt)
    f(k, u, p, t+c2end[i]*dt)
    integrator.destats.nf += 1
    @.. thread=thread tmp = tmp + δ2end[i]*u
    @.. thread=thread u   = γ12end[i]*u + γ22end[i]*tmp + γ32end[i]*uprev + β2end[i]*dt*k
    if integrator.opts.adaptive
      @.. thread=thread utilde = utilde + bhat2end[i]*dt*k
    end
  end

  stage_limiter!(u, integrator, p, t+dt)
  step_limiter!(u, integrator, p, t+dt)

  if integrator.opts.adaptive
    calculate_residuals!(atmp, utilde, uprev, u, integrator.opts.abstol, integrator.opts.reltol,integrator.opts.internalnorm,t, thread)
    integrator.EEst = integrator.opts.internalnorm(atmp,t)
  end
end

# 3S+ FSAL low storage methods: 3S methods adding another memory location for the embedded method (FSAL version)
function initialize!(integrator,cache::LowStorageRK3SpFSALConstantCache)
  integrator.fsalfirst = integrator.f(integrator.uprev, integrator.p, integrator.t) # Pre-start fsal
  integrator.destats.nf += 1
  integrator.kshortsize = 2
  integrator.k = typeof(integrator.k)(undef, integrator.kshortsize)

  # Avoid undefined entries if k is an array of arrays
  integrator.fsallast = zero(integrator.fsalfirst)
  integrator.k[1] = integrator.fsalfirst
  integrator.k[2] = integrator.fsallast
end

@muladd function perform_step!(integrator,cache::LowStorageRK3SpFSALConstantCache,repeat_step=false)
  @unpack t,dt,uprev,u,f,p = integrator
  @unpack γ12end, γ22end, γ32end, δ2end, β1, β2end, c2end, bhat1, bhat2end, bhatfsal = cache

  # u1
  tmp = uprev
  u   = tmp + β1*dt*integrator.fsalfirst
  if integrator.opts.adaptive
    utilde = bhat1*dt*integrator.fsalfirst
  end

  # other stages
  for i in eachindex(γ12end)
    k   = f(u, p, t+c2end[i]*dt)
    integrator.destats.nf += 1
    tmp = tmp + δ2end[i]*u
    u   = γ12end[i]*u + γ22end[i]*tmp + γ32end[i]*uprev + β2end[i]*dt*k
    if integrator.opts.adaptive
      utilde = utilde + bhat2end[i]*dt*k
    end
  end

  # FSAL
  integrator.fsallast = f(u, p, t+dt)
  integrator.destats.nf += 1

  if integrator.opts.adaptive
    utilde = utilde + bhatfsal*dt*integrator.fsallast
    atmp = calculate_residuals(utilde, uprev, u, integrator.opts.abstol, integrator.opts.reltol,integrator.opts.internalnorm,t)
    integrator.EEst = integrator.opts.internalnorm(atmp,t)
  end

  integrator.k[1] = integrator.fsalfirst
  integrator.k[2] = integrator.fsallast
  integrator.u = u
end

function initialize!(integrator,cache::LowStorageRK3SpFSALCache)
  @unpack k,fsalfirst = cache
  integrator.fsalfirst = fsalfirst
  integrator.fsallast = k
  integrator.kshortsize = 2
  resize!(integrator.k, integrator.kshortsize)
  integrator.k[1] = integrator.fsalfirst
  integrator.k[2] = integrator.fsallast
  integrator.f(integrator.fsalfirst,integrator.uprev,integrator.p,integrator.t) # FSAL for interpolation
  integrator.destats.nf += 1
end

@muladd function perform_step!(integrator,cache::LowStorageRK3SpFSALCache,repeat_step=false)
  @unpack t,dt,uprev,u,f,p = integrator
  @unpack k,tmp,utilde,atmp,stage_limiter!,step_limiter!,thread = cache
  @unpack γ12end, γ22end, γ32end, δ2end, β1, β2end, c2end, bhat1, bhat2end, bhatfsal = cache.tab

  # u1
  @.. thread=thread tmp = uprev
  @.. thread=thread u   = tmp + β1*dt*integrator.fsalfirst
  if integrator.opts.adaptive
    @.. thread=thread utilde = bhat1*dt*integrator.fsalfirst
  end

  # other stages
  for i in eachindex(γ12end)
    stage_limiter!(u, integrator, p, t+c2end[i]*dt)
    f(k, u, p, t+c2end[i]*dt)
    integrator.destats.nf += 1
    @.. thread=thread tmp = tmp + δ2end[i]*u
    @.. thread=thread u   = γ12end[i]*u + γ22end[i]*tmp + γ32end[i]*uprev + β2end[i]*dt*k
    if integrator.opts.adaptive
      @.. thread=thread utilde = utilde + bhat2end[i]*dt*k
    end
  end

  stage_limiter!(u, integrator, p, t+dt)
  step_limiter!(u, integrator, p, t+dt)

  # FSAL
  f(k, u, p, t+dt)
  integrator.destats.nf += 1

  if integrator.opts.adaptive
    @.. thread=thread utilde = utilde + bhatfsal*dt*k
    calculate_residuals!(atmp, utilde, uprev, u, integrator.opts.abstol, integrator.opts.reltol,integrator.opts.internalnorm,t, thread)
    integrator.EEst = integrator.opts.internalnorm(atmp,t)
  end
end

# 2R+ low storage methods
function initialize!(integrator,cache::LowStorageRK2RPConstantCache)
  integrator.fsalfirst = integrator.f(integrator.uprev, integrator.p, integrator.t) # Pre-start fsal
  integrator.destats.nf += 1
  integrator.kshortsize = 1
  integrator.k = typeof(integrator.k)(undef, integrator.kshortsize)

  # Avoid undefined entries if k is an array of arrays
  integrator.fsallast = zero(integrator.fsalfirst)
  integrator.k[1] = integrator.fsalfirst
end

@muladd function perform_step!(integrator,cache::LowStorageRK2RPConstantCache,repeat_step=false)
  @unpack t,dt,u,uprev,f,fsalfirst,p = integrator
  @unpack Aᵢ,Bₗ,B̂ₗ,Bᵢ,B̂ᵢ,Cᵢ = cache

  k   = fsalfirst
  integrator.opts.adaptive && (tmp = zero(uprev))

  #stages 1 to s-1
  for i in eachindex(Aᵢ)
    integrator.opts.adaptive && (tmp = tmp + (Bᵢ[i] - B̂ᵢ[i])*dt*k)
    gprev = u + Aᵢ[i]*dt*k
    u = u + Bᵢ[i]*dt*k
    k = f(gprev, p, t + Cᵢ[i]*dt)
    integrator.destats.nf += 1
  end

  #last stage
  integrator.opts.adaptive && (tmp = tmp + (Bₗ - B̂ₗ)*dt*k)
  u   = u  + Bₗ*dt*k

  #Error estimate
  if integrator.opts.adaptive
    atmp = calculate_residuals(tmp, uprev, u, integrator.opts.abstol, integrator.opts.reltol, integrator.opts.internalnorm, t)
    integrator.EEst = integrator.opts.internalnorm(atmp,t)
  end

  integrator.k[1] = integrator.fsalfirst
  integrator.fsallast = f(u, p, t+dt) # For interpolation, then FSAL'd
  integrator.destats.nf += 1
  integrator.u = u
end

function initialize!(integrator,cache::LowStorageRK2RPCache)
  @unpack k,fsalfirst = cache
  integrator.fsalfirst = fsalfirst
  integrator.fsallast = k
  integrator.kshortsize = 1
  resize!(integrator.k, integrator.kshortsize)
  integrator.k[1] = integrator.fsalfirst
  integrator.f(integrator.fsalfirst,integrator.uprev,integrator.p,integrator.t)
  integrator.destats.nf += 1
end

@muladd function perform_step!(integrator,cache::LowStorageRK2RPCache,repeat_step=false)
  @unpack t,dt,u,uprev,f,fsalfirst,p = integrator
  @unpack k,gprev,tmp,atmp = cache
  @unpack Aᵢ,Bₗ,B̂ₗ,Bᵢ,B̂ᵢ,Cᵢ = cache.tab

  @.. k   = fsalfirst
  integrator.opts.adaptive && (@.. tmp = zero(uprev))

  #stages 1 to s-1
  for i in eachindex(Aᵢ)
    integrator.opts.adaptive && (@.. tmp = tmp + (Bᵢ[i] - B̂ᵢ[i])*dt*k)
    @.. gprev = u + Aᵢ[i]*dt*k
    @.. u     = u + Bᵢ[i]*dt*k
    f(k, gprev, p, t + Cᵢ[i]*dt)
    integrator.destats.nf += 1
  end

  #last stage
  integrator.opts.adaptive && (@.. tmp = tmp + (Bₗ - B̂ₗ)*dt*k)
  @.. u   = u  + Bₗ*dt*k

  #Error estimate
  if integrator.opts.adaptive
    calculate_residuals!(atmp,tmp, uprev, u, integrator.opts.abstol, integrator.opts.reltol, integrator.opts.internalnorm, t)
    integrator.EEst = integrator.opts.internalnorm(atmp,t)
  end

  f(k, u, p, t+dt)
  integrator.destats.nf += 1
end

# 3R+ low storage methods
function initialize!(integrator,cache::LowStorageRK3RPConstantCache)
  integrator.fsalfirst = integrator.f(integrator.uprev, integrator.p, integrator.t) # Pre-start fsal
  integrator.destats.nf += 1
  integrator.kshortsize = 1
  integrator.k = typeof(integrator.k)(undef, integrator.kshortsize)

  # Avoid undefined entries if k is an array of arrays
  integrator.fsallast = zero(integrator.fsalfirst)
  integrator.k[1] = integrator.fsalfirst
end

@muladd function perform_step!(integrator,cache::LowStorageRK3RPConstantCache,repeat_step=false)
  @unpack t,dt,u,uprev,f,fsalfirst,p = integrator
  @unpack Aᵢ₁,Aᵢ₂,Bₗ,B̂ₗ,Bᵢ,B̂ᵢ,Cᵢ = cache

  fᵢ₋₂  = zero(fsalfirst)
  k     = fsalfirst
  uᵢ₋₁  = uprev
  uᵢ₋₂  = uprev
  integrator.opts.adaptive && (tmp = zero(uprev))

  #stages 1 to s-1
  for i in eachindex(Aᵢ₁)
    integrator.opts.adaptive && (tmp = tmp + (Bᵢ[i] - B̂ᵢ[i])*dt*k)
    gprev = uᵢ₋₂ + (Aᵢ₁[i]*k+Aᵢ₂[i]*fᵢ₋₂)*dt
    u     = u + Bᵢ[i]*dt*k
    fᵢ₋₂  = k
    uᵢ₋₂  = uᵢ₋₁
    uᵢ₋₁  = u
    k     = f(gprev, p, t + Cᵢ[i]*dt)
    integrator.destats.nf += 1
  end

  #last stage
  integrator.opts.adaptive && (tmp = tmp + (Bₗ - B̂ₗ)*dt*k)
  u   = u  + Bₗ*dt*k

  #Error estimate
  if integrator.opts.adaptive
    atmp = calculate_residuals(tmp, uprev, u, integrator.opts.abstol, integrator.opts.reltol, integrator.opts.internalnorm, t)
    integrator.EEst = integrator.opts.internalnorm(atmp,t)
  end

  integrator.k[1] = integrator.fsalfirst
  integrator.fsallast = f(u, p, t+dt) # For interpolation, then FSAL'd
  integrator.destats.nf += 1
  integrator.u = u
end

function initialize!(integrator,cache::LowStorageRK3RPCache)
  @unpack k,fsalfirst = cache
  integrator.fsalfirst = fsalfirst
  integrator.fsallast = k
  integrator.kshortsize = 1
  resize!(integrator.k, integrator.kshortsize)
  integrator.k[1] = integrator.fsalfirst
  integrator.f(integrator.fsalfirst,integrator.uprev,integrator.p,integrator.t)
  integrator.destats.nf += 1
end

@muladd function perform_step!(integrator,cache::LowStorageRK3RPCache,repeat_step=false)
  @unpack t,dt,u,uprev,f,fsalfirst,p = integrator
  @unpack k,uᵢ₋₁,uᵢ₋₂,gprev,fᵢ₋₂,tmp,atmp = cache
  @unpack Aᵢ₁,Aᵢ₂,Bₗ,B̂ₗ,Bᵢ,B̂ᵢ,Cᵢ = cache.tab

  @.. fᵢ₋₂  = zero(fsalfirst)
  @.. k     = fsalfirst
  integrator.opts.adaptive && (@.. tmp = zero(uprev))
  @.. uᵢ₋₁  = uprev
  @.. uᵢ₋₂  = uprev

  #stages 1 to s-1
  for i in eachindex(Aᵢ₁)
    integrator.opts.adaptive && (@.. tmp = tmp + (Bᵢ[i] - B̂ᵢ[i])*dt*k)
    @.. gprev = uᵢ₋₂ + (Aᵢ₁[i]*k+Aᵢ₂[i]*fᵢ₋₂)*dt
    @.. u     = u + Bᵢ[i]*dt*k
    @.. fᵢ₋₂  = k
    @.. uᵢ₋₂  = uᵢ₋₁
    @.. uᵢ₋₁  = u
    f(k, gprev, p, t + Cᵢ[i]*dt)
    integrator.destats.nf += 1
  end

  #last stage
  integrator.opts.adaptive && (@.. tmp = tmp + (Bₗ - B̂ₗ)*dt*k)
  @.. u   = u  + Bₗ*dt*k

  #Error estimate
  if integrator.opts.adaptive
    calculate_residuals!(atmp,tmp, uprev, u, integrator.opts.abstol, integrator.opts.reltol, integrator.opts.internalnorm, t)
    integrator.EEst = integrator.opts.internalnorm(atmp,t)
  end

  f(k, u, p, t+dt)
  integrator.destats.nf += 1
end

# 4R+ low storage methods
function initialize!(integrator,cache::LowStorageRK4RPConstantCache)
  integrator.fsalfirst = integrator.f(integrator.uprev, integrator.p, integrator.t) # Pre-start fsal
  integrator.destats.nf += 1
  integrator.kshortsize = 1
  integrator.k = typeof(integrator.k)(undef, integrator.kshortsize)

  # Avoid undefined entries if k is an array of arrays
  integrator.fsallast = zero(integrator.fsalfirst)
  integrator.k[1] = integrator.fsalfirst
end

@muladd function perform_step!(integrator,cache::LowStorageRK4RPConstantCache,repeat_step=false)
  @unpack t,dt,u,uprev,f,fsalfirst,p = integrator
  @unpack Aᵢ₁,Aᵢ₂,Aᵢ₃,Bₗ,B̂ₗ,Bᵢ,B̂ᵢ,Cᵢ = cache

  fᵢ₋₂  = zero(fsalfirst)
  fᵢ₋₃  = zero(fsalfirst)
  k     = fsalfirst
  uᵢ₋₁  = uprev
  uᵢ₋₂  = uprev
  uᵢ₋₃  = uprev
  integrator.opts.adaptive && (tmp = zero(uprev))

  #stages 1 to s-1
  for i in eachindex(Aᵢ₁)
    integrator.opts.adaptive && (tmp = tmp + (Bᵢ[i] - B̂ᵢ[i])*dt*k)
    gprev = uᵢ₋₃ + (Aᵢ₁[i]*k+Aᵢ₂[i]*fᵢ₋₂+Aᵢ₃[i]*fᵢ₋₃)*dt
    u     = u + Bᵢ[i]*dt*k
    fᵢ₋₃  = fᵢ₋₂
    fᵢ₋₂  = k
    uᵢ₋₃  = uᵢ₋₂
    uᵢ₋₂  = uᵢ₋₁
    uᵢ₋₁  = u
    k     = f(gprev, p, t + Cᵢ[i]*dt)
    integrator.destats.nf += 1
  end

  #last stage
  integrator.opts.adaptive && (tmp = tmp + (Bₗ - B̂ₗ)*dt*k)
  u   = u  + Bₗ*dt*k

  #Error estimate
  if integrator.opts.adaptive
    atmp = calculate_residuals(tmp, uprev, u, integrator.opts.abstol, integrator.opts.reltol, integrator.opts.internalnorm, t)
    integrator.EEst = integrator.opts.internalnorm(atmp,t)
  end

  integrator.k[1] = integrator.fsalfirst
  integrator.fsallast = f(u, p, t+dt) # For interpolation, then FSAL'd
  integrator.destats.nf += 1
  integrator.u = u
end

function initialize!(integrator,cache::LowStorageRK4RPCache)
  @unpack k,fsalfirst = cache
  integrator.fsalfirst = fsalfirst
  integrator.fsallast = k
  integrator.kshortsize = 1
  resize!(integrator.k, integrator.kshortsize)
  integrator.k[1] = integrator.fsalfirst
  integrator.f(integrator.fsalfirst,integrator.uprev,integrator.p,integrator.t)
  integrator.destats.nf += 1
end

@muladd function perform_step!(integrator,cache::LowStorageRK4RPCache,repeat_step=false)
  @unpack t,dt,u,uprev,f,fsalfirst,p = integrator
  @unpack k,uᵢ₋₁,uᵢ₋₂,uᵢ₋₃,gprev,fᵢ₋₂,fᵢ₋₃,tmp,atmp = cache
  @unpack Aᵢ₁,Aᵢ₂,Aᵢ₃,Bₗ,B̂ₗ,Bᵢ,B̂ᵢ,Cᵢ = cache.tab

  @.. fᵢ₋₂  = zero(fsalfirst)
  @.. fᵢ₋₃  = zero(fsalfirst)
  @.. k     = fsalfirst
  integrator.opts.adaptive && (@.. tmp = zero(uprev))
  @.. uᵢ₋₁  = uprev
  @.. uᵢ₋₂  = uprev
  @.. uᵢ₋₃  = uprev

  #stages 1 to s-1
  for i in eachindex(Aᵢ₁)
    integrator.opts.adaptive && (@.. tmp = tmp + (Bᵢ[i] - B̂ᵢ[i])*dt*k)
    @.. gprev = uᵢ₋₃ + (Aᵢ₁[i]*k+Aᵢ₂[i]*fᵢ₋₂+Aᵢ₃[i]*fᵢ₋₃)*dt
    @.. u     = u + Bᵢ[i]*dt*k
    @.. fᵢ₋₃  = fᵢ₋₂
    @.. fᵢ₋₂  = k
    @.. uᵢ₋₃  = uᵢ₋₂
    @.. uᵢ₋₂  = uᵢ₋₁
    @.. uᵢ₋₁  = u
    f(k, gprev, p, t + Cᵢ[i]*dt)
    integrator.destats.nf += 1
  end

  #last stage
  integrator.opts.adaptive && (@.. tmp = tmp + (Bₗ - B̂ₗ)*dt*k)
  @.. u   = u  + Bₗ*dt*k

  #Error estimate
  if integrator.opts.adaptive
    calculate_residuals!(atmp,tmp, uprev, u, integrator.opts.abstol, integrator.opts.reltol, integrator.opts.internalnorm, t)
    integrator.EEst = integrator.opts.internalnorm(atmp,t)
  end

  f(k, u, p, t+dt)
  integrator.destats.nf += 1
end

# 5R+ low storage methods
function initialize!(integrator,cache::LowStorageRK5RPConstantCache)
  integrator.fsalfirst = integrator.f(integrator.uprev, integrator.p, integrator.t) # Pre-start fsal
  integrator.destats.nf += 1
  integrator.kshortsize = 1
  integrator.k = typeof(integrator.k)(undef, integrator.kshortsize)

  # Avoid undefined entries if k is an array of arrays
  integrator.fsallast = zero(integrator.fsalfirst)
  integrator.k[1] = integrator.fsalfirst
end

@muladd function perform_step!(integrator,cache::LowStorageRK5RPConstantCache,repeat_step=false)
  @unpack t,dt,u,uprev,f,fsalfirst,p = integrator
  @unpack Aᵢ₁,Aᵢ₂,Aᵢ₃,Aᵢ₄,Bₗ,B̂ₗ,Bᵢ,B̂ᵢ,Cᵢ = cache

  fᵢ₋₂  = zero(fsalfirst)
  fᵢ₋₃  = zero(fsalfirst)
  fᵢ₋₄  = zero(fsalfirst)
  k     = fsalfirst
  uᵢ₋₁  = uprev
  uᵢ₋₂  = uprev
  uᵢ₋₃  = uprev
  uᵢ₋₄  = uprev
  integrator.opts.adaptive && (tmp = zero(uprev))

  #stages 1 to s-1
  for i in eachindex(Aᵢ₁)
    integrator.opts.adaptive && (tmp = tmp + (Bᵢ[i] - B̂ᵢ[i])*dt*k)
    gprev = uᵢ₋₄ + (Aᵢ₁[i]*k+Aᵢ₂[i]*fᵢ₋₂+Aᵢ₃[i]*fᵢ₋₃+Aᵢ₄[i]*fᵢ₋₄)*dt
    u     = u + Bᵢ[i]*dt*k
    fᵢ₋₄  = fᵢ₋₃
    fᵢ₋₃  = fᵢ₋₂
    fᵢ₋₂  = k
    uᵢ₋₄  = uᵢ₋₃
    uᵢ₋₃  = uᵢ₋₂
    uᵢ₋₂  = uᵢ₋₁
    uᵢ₋₁  = u
    k     = f(gprev, p, t + Cᵢ[i]*dt)
    integrator.destats.nf += 1
  end

  #last stage
  integrator.opts.adaptive && (tmp = tmp + (Bₗ - B̂ₗ)*dt*k)
  u   = u  + Bₗ*dt*k

  #Error estimate
  if integrator.opts.adaptive
    atmp = calculate_residuals(tmp, uprev, u, integrator.opts.abstol, integrator.opts.reltol, integrator.opts.internalnorm, t)
    integrator.EEst = integrator.opts.internalnorm(atmp,t)
  end

  integrator.k[1] = integrator.fsalfirst
  integrator.fsallast = f(u, p, t+dt) # For interpolation, then FSAL'd
  integrator.destats.nf += 1
  integrator.u = u
end

function initialize!(integrator,cache::LowStorageRK5RPCache)
  @unpack k,fsalfirst = cache
  integrator.fsalfirst = fsalfirst
  integrator.fsallast = k
  integrator.kshortsize = 1
  resize!(integrator.k, integrator.kshortsize)
  integrator.k[1] = integrator.fsalfirst
  integrator.f(integrator.fsalfirst,integrator.uprev,integrator.p,integrator.t)
  integrator.destats.nf += 1
end

@muladd function perform_step!(integrator,cache::LowStorageRK5RPCache,repeat_step=false)
  @unpack t,dt,u,uprev,f,fsalfirst,p = integrator
  @unpack k,uᵢ₋₁,uᵢ₋₂,uᵢ₋₃,uᵢ₋₄,gprev,fᵢ₋₂,fᵢ₋₃,fᵢ₋₄,tmp,atmp = cache
  @unpack Aᵢ₁,Aᵢ₂,Aᵢ₃,Aᵢ₄,Bₗ,B̂ₗ,Bᵢ,B̂ᵢ,Cᵢ = cache.tab

  @.. fᵢ₋₂  = zero(fsalfirst)
  @.. fᵢ₋₃  = zero(fsalfirst)
  @.. fᵢ₋₄  = zero(fsalfirst)
  @.. k     = fsalfirst
  integrator.opts.adaptive && (@.. tmp = zero(uprev))
  @.. uᵢ₋₁  = uprev
  @.. uᵢ₋₂  = uprev
  @.. uᵢ₋₃  = uprev
  @.. uᵢ₋₄  = uprev

  #stages 1 to s-1
  for i in eachindex(Aᵢ₁)
    integrator.opts.adaptive && (@.. tmp = tmp + (Bᵢ[i] - B̂ᵢ[i])*dt*k)
    @.. gprev = uᵢ₋₄ + (Aᵢ₁[i]*k+Aᵢ₂[i]*fᵢ₋₂+Aᵢ₃[i]*fᵢ₋₃+Aᵢ₄[i]*fᵢ₋₄)*dt
    @.. u     = u + Bᵢ[i]*dt*k
    @.. fᵢ₋₄  = fᵢ₋₃
    @.. fᵢ₋₃  = fᵢ₋₂
    @.. fᵢ₋₂  = k
    @.. uᵢ₋₄  = uᵢ₋₃
    @.. uᵢ₋₃  = uᵢ₋₂
    @.. uᵢ₋₂  = uᵢ₋₁
    @.. uᵢ₋₁  = u
    f(k, gprev, p, t + Cᵢ[i]*dt)
    integrator.destats.nf += 1
  end

  #last stage
  integrator.opts.adaptive && (@.. tmp = tmp + (Bₗ - B̂ₗ)*dt*k)
  @.. u   = u  + Bₗ*dt*k

  #Error estimate
  if integrator.opts.adaptive
    calculate_residuals!(atmp,tmp, uprev, u, integrator.opts.abstol, integrator.opts.reltol, integrator.opts.internalnorm, t)
    integrator.EEst = integrator.opts.internalnorm(atmp,t)
  end

  f(k, u, p, t+dt)
  integrator.destats.nf += 1
end


text:
# Test linearity of sources
# Author: Philipp Witte, pwitte@eos.ubc.ca
# Date: January 2017
#

using PyCall, PyPlot, opesciSLIM.TimeModeling

## Set up model structure
n = (120,100)	# (x,y,z) or (x,z)
d = (10.,10.)
o = (0.,0.)

# Velocity [km/s]
v = ones(Float32,n) + 0.4f0
v[:,Int(round(end/2)):end] = 4f0

# Slowness squared [s^2/km^2]
m = (1f0./v).^2

# Setup info and model structure
nsrc = 1
ntComp = 250
info = Info(prod(n), nsrc, ntComp)	# number of gridpoints, number of experiments, number of computational time steps
model = Model(n,d,o,m)

## Set up receiver geometry
nxrec = 120
xrec = linspace(50f0,1150f0,nxrec)
yrec = 0f0
zrec = linspace(50f0,50f0,nxrec)

# receiver sampling and recording time
timeR = 1000f0	# receiver recording time [ms]
dtR = 4f0	# receiver sampling interval

# Set up receiver structure
recGeometry = Geometry(xrec,yrec,zrec;dt=dtR,t=timeR,nsrc=nsrc)

## Set up source geometries (cell array with source locations for each shot)
xsrc1 = 300f0
ysrc1 = 0f0
zsrc1 = 20f0

xsrc2 = 600f0
ysrc2 = 0f0
zsrc2 = 20f0

# source sampling and number of time steps
timeS = 1000f0
dtS = 4f0	# receiver sampling interval

# Set up source structure
srcGeometry1 = Geometry(xsrc1,ysrc1,zsrc1;dt=dtS,t=timeS)
srcGeometry2 = Geometry(xsrc2,ysrc2,zsrc2;dt=dtS,t=timeS)

# setup wavelet
f0 = 0.01f0
wavelet = source(timeS,dtS,f0)

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

# Modeling operators
Pr = joProjection(info,recGeometry)
Ps1 = joProjection(info,srcGeometry1)
Ps2 = joProjection(info,srcGeometry2)
F = joModeling(info,model)
q1 = joData(srcGeometry1,wavelet)
q2 = joData(srcGeometry2,wavelet)

d1 = Pr*F*Ps1'*q1
d2 = Pr*F*Ps2'*q2
d3 = Pr*F*(Ps1'*q1 + Ps2'*q2)
d4 = Pr*F*(Ps1'*q1 - Ps2'*q2)

println("addition: ", norm(d3 - (d1 + d2)))
println("subtraction: ", norm(d4 - (d1 - d2)))

max_stars_repo_name: SINBADconsortium/opesciSLIM

max_issues_repo_name: SINBADconsortium/opesciSLIM

max_forks_repo_name: SINBADconsortium/opesciSLIM


text:
module TWPBVP_Test1
	using TWPBVP
	const ϵ = 1e-2

	function f(x, y, dy)
		println("f: input x = $x, y = $y")
		dy[1] = y[2]
		dy[2] = 1/ϵ*(-exp(y[1]) + 0.5π*sin(0.5π*x)*exp(2y[1]))
	end

	function df(x, y, jac)
		jac[1,1] = 0.0
		jac[1,2] = 1.0
		jac[2,1] = 1/ϵ*(-exp(y[1])+π*sin(0.5π*x)*exp(2y[1]))
		jac[1,2] = 0.0
	end

	function g(i, y)
		y[1]
	end

	function dg(i, y, jac)
		jac[1] = 1.0
		jac[2] = 0.0
	end

	function perform()

		# Other constants & values
		fxpoint = Nullable{Vector{Float64}}()
		local ltol :: Vector{FInt} = [1]
		tol = [1e-2]
		u = Array(Float64, 2, 1000)
		xx = Array(Float64, 1000)
		wrk = Array(Float64, 10000)
		iwrk = Array(FInt, 6000)
		local num_left_bc :: FInt = 1
		left = 0.0
		right = 1.0
		linear = false
		giveu = false
		givmsh = false
		nmesh = Ref{FInt}(0)
		nmax = Ref{FInt}(0)
		κ1 = Ref(1.0)
		κ = Ref(1.0)
		γ1 = Ref(1.0)
		result = Ref{FInt}(0)

		twpbvpc(num_left_bc, left, right, fxpoint, ltol, tol, linear, givmsh, giveu, nmesh, xx, u, nmax, wrk, iwrk, f, df, g, dg, κ1, γ1, κ, result)
	end

end

module TWPBVP_Test2

	using TWPBVP

	function f(t,y, dy)
		dy[1] = y[2]
		dy[2] = -y[1]
	end

	function df(t, y, jac)
		jac[1,1] = 0.0
		jac[1,2] = 1.0
		jac[2,1] = -1.0
		jac[2,2] = 0.0
	end

	g(i, y) = i == 1 ? y[1] - 1.0 : y[1] + 1.0

	function dg(i :: FInt, y :: Vector{Float64}, grad :: Vector{Float64})
		grad[:] = [1.0, 0.0]
	end

	function perform()

		fixpoints = Nullable{Vector{Float64}}()
		local ltol :: Vector{FInt} = [2]
		tol = [1e-6]
		u = Array(Float64, 2, 1000)
		xx = Array(Float64, 1000)

		wrk = Array(Float64, 10000)
		iwrk = Array(FInt, 6000)
		local num_left_bc :: FInt = 1
		left = 0.0
		right = float(π)
		linear = false
		giveu = true

		givmsh = true

		local nmesh0 :: FInt = 20

		nmesh = Ref{FInt}(nmesh0)

		nmax = Ref{FInt}(0)
		κ1 = Ref(1.0)
		κ = Ref(1.0)
		γ1 = Ref(1.0)
		result = Ref{FInt}(-100)

		xx[1:nmesh0] = linspace(0.0, π, nmesh0)

		for i=1:nmesh0
			u[:, i] = [-2/π*xx[i] + 1, -2.0/π]
		end
		println("xx0 = $(xx[1:nmesh0+2])")

		twpbvpc(num_left_bc, left, right, fixpoints, ltol, tol, linear, givmsh, giveu, nmesh, xx, u, nmax, wrk, iwrk, f, df, g, dg, κ1, γ1, κ, result)

		println("Result is $(result[])")
		println("nmesh = $(nmesh[])")

	end

end

using Base.Test
import TWPBVP_Test1

TT = TWPBVP_Test1

TT.perform()


text:
using GAIO

# the Lorenz system
const σ, ρ, β = 10.0, 28.0, 0.4
v((x,y,z)) = (σ*(y-x), ρ*x-y-x*z, x*y-β*z)
f(x) = rk4_flow_map(v, x)

center, radius = (0,0,25), (30,30,30)
P = BoxPartition(Box(center, radius), (128,128,128))
F = BoxMap(f, P, no_of_points=200)

x = (sqrt(β*(ρ-1)), sqrt(β*(ρ-1)), ρ-1)         # equilibrium
W = unstable_set!(F, P[x])

T = TransferOperator(F, W)
(λ, ev) = eigs(T)

plot(log∘abs∘ev[1])