"""
    barsandstripes(nsamples, nvariables)
Generates a test data set. To see the structure in the data set, run e. g.
`reshape(barsandstripes(1, 16), 4, 4)` a few times.

Example from:
MacKay, D. (2003). Information Theory, Inference, and Learning Algorithms
"""
function barsandstripes(nsamples::Int, nvariables::Int)
   squareside = sqrt(nvariables)
   if squareside != floor(squareside)
      error("Number of variables must be a square number")
   end
   squareside = round(Int, squareside)
   x = zeros(nsamples, nvariables)
   for i in 1:nsamples
      sample = hcat([(rand() < 0.5 ? ones(squareside) : zeros(squareside))
            for j in 1:squareside]...)
      if rand() < 0.5
         sample = transpose(sample)
      end
      x[i,:] .= sample[:]
      fill!(sample, 0.0)
   end
   x
end

"""
    batchparallelized(f, n, op)
Distributes the work for executing the function `f` `n` times
on all the available workers and reduces the results with the operator `op`.
`f` is a function that gets a number (of tasks) to execute the tasks.

# Example:
    batchparallelized(n -> aislogimpweights(dbm; nparticles = n), 100, vcat)
"""
function batchparallelized(f::Function, n::Int, op::Function)
   batches = mostevenbatches(n)
   if length(batches) > 1
      return @sync @distributed (op) for batch in batches
         f(batch)
      end
   else
      return f(n)
   end
end

"""
    crossvalidation(x, monitoredfit; ...)
Performs k-fold cross-validation, given
* the data set `x` and
* `monitoredfit`: a function that fits and evaluates a model.
  As arguments it must accept:
  - a training data data set
  - a `DataDict` containing the evaluation data.

The return values of the calls to the `monitoredfit` function
are concatenated with `vcat`.
If the monitoredfit function returns `Monitor` objects,
`crossvalidation` returns a combined `Monitor` object that can be displayed
by creating a cross-validation plot via
`BoltzmannMachinesPlots.crossvalidationplot`.

# Optional named argument:
* `kfold`: specifies the `k` in "`k`-fold" (defaults to 10).

    crossvalidation(x, monitoredfit, pars; ...)
If additionaly a vector of parameters `pars` is given, `monitoredfit`
also expects an additional parameter from the parameter set.
"""
function crossvalidation(x::Matrix{Float64}, monitoredfit::Function, pars...;
      kfold::Int = 10)

   vcat(pmap(monitoredfit, crossvalidationargs(x, pars...)...)...)
end

"""
    crossvalidationargs(x, pars...; )

Returns a tuple of argument vectors containing the parameters for a function
such as the `monitoredfit` argument in `crossvalidation`.

Usage example:
    map(monitoredfit, crossvalidationargs(x)...)

# Optional named argument:
* `kfold`: see `crossvalidation`.
"""
function crossvalidationargs(x::Matrix{Float64}, pars ...; kfold::Int = 10)

   nsamples = size(x, 1)
   batchranges = ranges(mostevenbatches(nsamples, kfold))

   args_data = Vector{Tuple{Matrix{Float64}, DataDict}}(undef, kfold)
   for i in eachindex(batchranges)
      rng = batchranges[i]
      trainingdata = x[[!(j in rng) for j in 1:nsamples], :]
      evaluationdata = x[rng, :]
      datadict = DataDict(string(i) => evaluationdata)
      args_data[i] = (trainingdata, datadict)
   end

   # this implementation is inefficient and overly complicated

   if isempty(pars)
      args = map(arg -> (arg,), args_data)
   else
      args = (args_data, map(p -> collect(p), pars)...)
      args = vec(collect(Iterators.product(args...)))
   end

   # vector of tuples to tuples of vectors
   (
      map(arg -> arg[1][1], args), # x
      map(arg -> arg[1][2], args), # datadict
      [map(arg -> arg[i], args) for i in 2:length(args[1])]...
   )
end

"""
A function with no arguments doing nothing.
Usable as default argument for functions as arguments.W
"""
function emptyfunc
end

"""
    log1pexp(x)
Calculates log(1+exp(x)). For sufficiently large values of x, the approximation
log(1+exp(x)) ≈ x is used. This is useful to prevent overflow.
"""
function log1pexp(x::Float64)
   if x > 34
      return x
   else
      return log1p(exp(x))
   end
end

function logit(p::Float64)
   -log.(1.0 ./ p .- 1.0)
end

"""
    mostevenbatches(ntasks)
    mostevenbatches(ntasks, nbatches)
Splits a number of tasks `ntasks` into a number of batches `nbatches`.
The number of batches is by default `min(nworkers(), ntasks)`.
The returned result is a vector containing the numbers of tasks for each batch.
"""
function mostevenbatches(ntasks::Int, nbatches::Int = min(nworkers(), ntasks))

   minparticlesperbatch, nbatcheswithplus1 = divrem(ntasks, nbatches)
   batches = Vector{Int}(undef, nbatches)
   for i = 1:nbatches
      if i <= nbatcheswithplus1
         batches[i] = minparticlesperbatch + 1
      else
         batches[i] = minparticlesperbatch
      end
   end
   batches
end

"""
    curvebundles(...)
Generates an example dataset that can be visualized as bundles of trend curves
with added noise. Additional binary columns with labels may be added.

## Optional named arguments:
* `nbundles`: number of bundles
* `nperbundle`: number of sequences per bundle
* `nvariables`: number of variables in the sequences
* `noisesd`: standard deviation of the noise added on all sequences
* `addlabels`: add leading columns to the resulting dataset, specifying the
   membership to a bundle
* `pbreak`: probability that an intermediate point in a sequence is a
   breakpoint, defaults to 0.2.
* `breakval`: a function that expects no input and generates a single
  (random) value for a defining point of a piecewise linear sequence.
  Defaults to `rand`.

# Example:
To quickly grasp the idea, plot generated samples against the variable index:

    x = BMs.curvebundles(nvariables = 10, nbundles = 3,
                       nperbundle = 4, noisesd = 0.03,
                       addlabels = true)
    BoltzmannMachinesPlots.plotcurvebundles(x)
"""
function curvebundles(;
      nbundles::Int = 3,
      nperbundle::Int = 50,
      nvariables::Int = 10,
      noisesd::Float64 = 0.05,
      addlabels::Bool = false,
      pbreak::Float64 = 0.2,
      breakval::Function = rand)

   x = piecewiselinearsequences(nbundles, nvariables;
         pbreak = pbreak, breakval = breakval)
   x = repeat(x, nperbundle)
   nsamples = size(x, 1)
   x .+= noisesd * randn(nsamples, nvariables)

   if addlabels
      nlabelvars = Int(ceil(log(nbundles)/log(2)))
      labels = Matrix{Float64}(nbundles, nlabelvars)
      labels[1,:] .= 0.0
      for j = 2:nbundles
         labels[j,:] .= labels[j-1,:]
         next!(view(labels, j, :))
      end
      labels = repeat(labels, nperbundle)
      x = hcat(labels, x)
   end

   x = x[randperm(nsamples), :]
   x
end

"""
    piecewiselinearsequences(nsequences, nvariables; ...)
Generates a dataset consisting of samples with values that
are piecewise linear functions of the variable index.

Optional named arguments: `pbreak`, `breakval`,
see `piecewiselinearsequencebundles`.
"""
function piecewiselinearsequences(nsequences::Int, nvariables::Int;
      pbreak::Float64 = 0.2, breakval::Function = rand)

   inbetweenvariables = 2:(nvariables-1)
   x = Matrix{Float64}(undef, nsequences, nvariables)
   for i = 1:nsequences
      breakpointindexes = [1; randsubseq(inbetweenvariables, pbreak); nvariables]
      breakpointvalues = [breakval() for i = 1:length(breakpointindexes)]
      lastbreakpoint = 1
      x[i, breakpointindexes] .= breakpointvalues
      for breakpoint in breakpointindexes[2:end]
         valdiff = (x[i, breakpoint] - x[i, lastbreakpoint]) /
               (breakpoint - lastbreakpoint)
         for j = (lastbreakpoint + 1):(breakpoint - 1)
            x[i, j] = x[i, j - 1] + valdiff
         end
         lastbreakpoint = breakpoint
      end
   end
   x
end

"""
    splitdata(x, ratio)
Splits the data set `x` randomly in two data sets `x1` and `x2`, such that
the ratio `n2`/`n1` of the numbers of lines/samples in `x1` and `x2` is
approximately equal to the given `ratio`.
"""
function splitdata(x::Matrix{Float64}, ratio::Float64)
   nsamples = size(x,1)
   testidxs = randperm(nsamples)[1:(round(Int, nsamples*ratio))]
   xtest = x[testidxs,:]
   trainidxs = setdiff(1:nsamples, testidxs)
   xtraining = x[setdiff(1:nsamples, testidxs),:]
   xtraining, xtest, trainidxs, testidxs
end

"""
    DBindex(x, cl; p, q)
Computes the Davies-Bouldin Index for cluster validation.
"""
function DBindex(x::Array{Float64,2}, cl::Array{Int64,1}; p::Int = 1, q::Int = 2)

   n = size(cl,1)
   k = sort(unique(cl))

   centers = Array{Float64,2}(undef,(size(k,1),size(x,2)))
   for i = 1:size(k,1)
      for j = 1:size(x,2)
         centers[i,j] =  mean(x[cl .== k[i], j])
      end
   end

   # S is the within-cluster dispersion
   S = zeros(size(k,1))
   for i = 1:size(k,1)
      ind = (cl .== k[i])
      if sum(ind) > 1
         centerI = centers[i,:]
         centerI = reshape(repeat(centerI', sum(ind)), (sum(ind), size(x,2)))
         tmp = (x[ind,:] .- centerI).^2
         S[i] = mean(sqrt.(mapslices(sum, tmp, dims = 2)).^q).^(1/q)
      else
         S[i] = 0.0
      end
   end

   # R is the between-cluster similarity as a function of the two within-cluster dispersions
   # M is the distance between cluster centroids
   M = pairwise(Euclidean(), centers, dims=1)
   R = Array{Float64,2}(undef,(size(k,1),size(k,1)))
   r = zeros(size(k,1))
   for i = 1:size(k,1)
      for j = 1:size(k,1)
         R[i,j] = (S[i] + S[j]) / M[i,j]
      end
      r[i] = maximum(R[i,:][isfinite.(R[i,:])])
   end

   DB = mean(r[isfinite.(r)])
   DB
end

# Number of parameters for Negative Binomial DBM

function nparams(dbm::Array{AbstractXBernoulliRBM,1})

   params = zeros()
   for i in 1:size(dbm,1)
      if  typeof(dbm[i]) == NegativeBinomialBernoulliRBM
         tmp = (size(dbm[i].weights,1) * size(dbm[i].weights,2)) + size(dbm[i].visbias,1) + size(dbm[i].hidbias,1) + size(dbm[i].inversedispersion,1)
         params .+= tmp
      else
         tmp = (size(dbm[i].weights,1) * size(dbm[i].weights,2)) + size(dbm[i].visbias,1) + size(dbm[i].hidbias,1)
         params .+= tmp
      end
   end
   params[1]
end


abstract type RosenbrockMutableCache <: OrdinaryDiffEqMutableCache end
################################################################################

# Shampine's Low-order Rosenbrocks

mutable struct Rosenbrock23Cache{uType,uArrayType,rateType,du2Type,LinuType,vecuType,JType,WType,TabType,TFType,UFType,F,JCType,GCType} <: RosenbrockMutableCache
  u::uType
  uprev::uType
  k₁::rateType
  k₂::rateType
  k₃::rateType
  du1::rateType
  du2::du2Type
  f₁::rateType
  vectmp::vecuType
  vectmp2::vecuType
  vectmp3::vecuType
  fsalfirst::rateType
  fsallast::rateType
  dT::uArrayType
  J::JType
  W::WType
  tmp::uArrayType
  tab::TabType
  tf::TFType
  uf::UFType
  linsolve_tmp::LinuType
  linsolve_tmp_vec::vecuType
  linsolve::F
  jac_config::JCType
  grad_config::GCType
end

u_cache(c::Rosenbrock23Cache) = (c.dT,c.tmp)
du_cache(c::Rosenbrock23Cache) = (c.k₁,c.k₂,c.k₃,c.du1,c.du2,c.f₁,c.fsalfirst,c.fsallast,c.linsolve_tmp)
jac_cache(c::Rosenbrock23Cache) = (c.J,c.W)
vecu_cache(c::Rosenbrock23Cache) = (c.vectmp,c.vectmp2,c.vectmp3)

mutable struct Rosenbrock32Cache{uType,uArrayType,rateType,du2Type,LinuType,vecuType,JType,WType,TabType,TFType,UFType,F,JCType,GCType} <: RosenbrockMutableCache
  u::uType
  uprev::uType
  k₁::rateType
  k₂::rateType
  k₃::rateType
  du1::rateType
  du2::du2Type
  f₁::rateType
  vectmp::vecuType
  vectmp2::vecuType
  vectmp3::vecuType
  fsalfirst::rateType
  fsallast::rateType
  dT::uArrayType
  J::JType
  W::WType
  tmp::uArrayType
  tab::TabType
  tf::TFType
  uf::UFType
  linsolve_tmp::LinuType
  linsolve_tmp_vec::vecuType
  linsolve::F
  jac_config::JCType
  grad_config::GCType
end

u_cache(c::Rosenbrock32Cache) = (c.dT,c.tmp)
du_cache(c::Rosenbrock32Cache) = (c.k₁,c.k₂,c.k₃,c.du1,c.du2,c.f₁,c.fsalfirst,c.fsallast,c.linsolve_tmp)
jac_cache(c::Rosenbrock32Cache) = (c.J,c.W)
vecu_cache(c::Rosenbrock32Cache) = (c.vectmp,c.vectmp2,c.vectmp3)

function alg_cache(alg::Rosenbrock23,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{true}})
  k₁ = zero(rate_prototype)
  k₂ = zero(rate_prototype)
  k₃ = zero(rate_prototype)
  du1 = zero(rate_prototype)
  du2 = zero(rate_prototype)
  # f₀ = similar(u) fsalfirst
  f₁ = zero(rate_prototype)
  vectmp = vec(similar(u,axes(u)))
  vectmp2 = vec(similar(u,axes(u)))
  vectmp3 = vec(similar(u,axes(u)))
  fsalfirst = zero(rate_prototype)
  fsallast = zero(rate_prototype)
  dT = similar(u,axes(u))
  if DiffEqBase.has_jac(f) && !DiffEqBase.has_invW(f) && f.jac_prototype != nothing
    W = WOperator(f, dt)
    J = nothing # is J = W.J better?
  else
    J = fill(zero(uEltypeNoUnits),length(u),length(u)) # uEltype?
    W = similar(J)
  end
  tmp = similar(u,axes(u))
  tab = Rosenbrock23ConstantCache(uEltypeNoUnits,identity,identity)
  tf = DiffEqDiffTools.TimeGradientWrapper(f,uprev,p)
  uf = DiffEqDiffTools.UJacobianWrapper(f,t,p)
  linsolve_tmp = similar(u,axes(u))
  linsolve_tmp_vec = vec(linsolve_tmp)
  linsolve = alg.linsolve(Val{:init},uf,u)

  grad_config = build_grad_config(alg,f,tf,du1,t)
  jac_config = build_jac_config(alg,f,uf,du1,uprev,u,tmp,du2)

  Rosenbrock23Cache(u,uprev,k₁,k₂,k₃,du1,du2,f₁,vectmp,vectmp2,vectmp3,
                    fsalfirst,fsallast,dT,J,W,tmp,tab,tf,uf,linsolve_tmp,
                    linsolve_tmp_vec,linsolve,jac_config,grad_config)
end

function alg_cache(alg::Rosenbrock32,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{true}})
  k₁ = zero(rate_prototype)
  k₂ = zero(rate_prototype)
  k₃ = zero(rate_prototype)
  du1 = zero(rate_prototype)
  du2 = zero(rate_prototype)
  # f₀ = similar(u) fsalfirst
  f₁ = zero(rate_prototype)
  vectmp = vec(similar(u,axes(u)))
  vectmp2 = vec(similar(u,axes(u)))
  vectmp3 = vec(similar(u,axes(u)))
  fsalfirst = zero(rate_prototype)
  fsallast = zero(rate_prototype)
  dT = similar(u,axes(u))
  if DiffEqBase.has_jac(f) && !DiffEqBase.has_invW(f) && f.jac_prototype != nothing
    W = WOperator(f, dt)
    J = nothing # is J = W.J better?
  else
    J = fill(zero(uEltypeNoUnits),length(u),length(u)) # uEltype?
    W = similar(J)
  end
  tmp = similar(u,axes(u))
  tab = Rosenbrock32ConstantCache(uEltypeNoUnits,identity,identity)

  tf = DiffEqDiffTools.TimeGradientWrapper(f,uprev,p)
  uf = DiffEqDiffTools.UJacobianWrapper(f,t,p)
  linsolve_tmp = similar(u,axes(u))
  linsolve_tmp_vec = vec(linsolve_tmp)
  linsolve = alg.linsolve(Val{:init},uf,u)
  grad_config = build_grad_config(alg,f,tf,du1,t)
  jac_config = build_jac_config(alg,f,uf,du1,uprev,u,tmp,du2)
  Rosenbrock32Cache(u,uprev,k₁,k₂,k₃,du1,du2,f₁,vectmp,vectmp2,vectmp3,fsalfirst,fsallast,dT,J,W,tmp,tab,tf,uf,linsolve_tmp,linsolve_tmp_vec,linsolve,jac_config,grad_config)
end

struct Rosenbrock23ConstantCache{T,TF,UF} <: OrdinaryDiffEqConstantCache
  c₃₂::T
  d::T
  tf::TF
  uf::UF
end

function Rosenbrock23ConstantCache(T::Type,tf,uf)
  c₃₂ = T(6 + sqrt(2))
  d = T(1/(2+sqrt(2)))
  Rosenbrock23ConstantCache(c₃₂,d,tf,uf)
end

function alg_cache(alg::Rosenbrock23,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{false}})
  tf = DiffEqDiffTools.TimeDerivativeWrapper(f,u,p)
  uf = DiffEqDiffTools.UDerivativeWrapper(f,t,p)
  Rosenbrock23ConstantCache(uEltypeNoUnits,tf,uf)
end

struct Rosenbrock32ConstantCache{T,TF,UF} <: OrdinaryDiffEqConstantCache
  c₃₂::T
  d::T
  tf::TF
  uf::UF
end

function Rosenbrock32ConstantCache(T::Type,tf,uf)
  c₃₂ = T(6 + sqrt(2))
  d = T(1/(2+sqrt(2)))
  Rosenbrock32ConstantCache(c₃₂,d,tf,uf)
end

function alg_cache(alg::Rosenbrock32,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{false}})
  tf = DiffEqDiffTools.TimeDerivativeWrapper(f,u,p)
  uf = DiffEqDiffTools.UDerivativeWrapper(f,t,p)
  Rosenbrock32ConstantCache(uEltypeNoUnits,tf,uf)
end

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

### 3rd order specialized Rosenbrocks

struct Rosenbrock33ConstantCache{TF,UF,Tab} <: OrdinaryDiffEqConstantCache
  tf::TF
  uf::UF
  tab::Tab
end

mutable struct Rosenbrock33Cache{uType,uArrayType,rateType,du2Type,LinuType,vecuType,JType,WType,TabType,TFType,UFType,F,JCType,GCType} <: RosenbrockMutableCache
  u::uType
  uprev::uType
  du::rateType
  du1::rateType
  du2::du2Type
  vectmp::vecuType
  vectmp2::vecuType
  vectmp3::vecuType
  vectmp4::vecuType
  fsalfirst::rateType
  fsallast::rateType
  dT::uArrayType
  J::JType
  W::WType
  tmp::uArrayType
  tab::TabType
  tf::TFType
  uf::UFType
  linsolve_tmp::LinuType
  linsolve_tmp_vec::vecuType
  linsolve::F
  jac_config::JCType
  grad_config::GCType
end

u_cache(c::Rosenbrock33Cache) = (c.dT,c.tmp)
du_cache(c::Rosenbrock33Cache) = (c.k₁,c.k₂,c.k₃,c.du1,c.du2,c.f₁,c.fsalfirst,c.fsallast,c.linsolve_tmp)
jac_cache(c::Rosenbrock33Cache) = (c.J,c.W)
vecu_cache(c::Rosenbrock33Cache) = (c.vectmp,c.vectmp2,c.vectmp3)

function alg_cache(alg::ROS3P,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{true}})
  du = zero(rate_prototype)
  du1 = zero(rate_prototype)
  du2 = zero(rate_prototype)
  vectmp = vec(similar(u,axes(u)))
  vectmp2 = vec(similar(u,axes(u)))
  vectmp3 = vec(similar(u,axes(u)))
  vectmp4 = vec(similar(u,axes(u)))
  fsalfirst = zero(rate_prototype)
  fsallast = zero(rate_prototype)
  dT = similar(u,axes(u))
  if DiffEqBase.has_jac(f) && !DiffEqBase.has_invW(f) && f.jac_prototype != nothing
    W = WOperator(f, dt)
    J = nothing # is J = W.J better?
  else
    J = fill(zero(uEltypeNoUnits),length(u),length(u)) # uEltype?
    W = similar(J)
  end
  tmp = similar(u,axes(u))
  tab = ROS3PConstantCache(real(uBottomEltypeNoUnits),real(tTypeNoUnits))
  tf = DiffEqDiffTools.TimeGradientWrapper(f,uprev,p)
  uf = DiffEqDiffTools.UJacobianWrapper(f,t,p)
  linsolve_tmp = similar(u,axes(u))
  linsolve_tmp_vec = vec(linsolve_tmp)
  linsolve = alg.linsolve(Val{:init},uf,u)
  grad_config = build_grad_config(alg,f,tf,du1,t)
  jac_config = build_jac_config(alg,f,uf,du1,uprev,u,tmp,du2)
  Rosenbrock33Cache(u,uprev,du,du1,du2,vectmp,vectmp2,vectmp3,vectmp4,
                    fsalfirst,fsallast,dT,J,W,tmp,tab,tf,uf,linsolve_tmp,
                    linsolve_tmp_vec,linsolve,jac_config,grad_config)
end

function alg_cache(alg::ROS3P,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{false}})
  tf = DiffEqDiffTools.TimeDerivativeWrapper(f,u,p)
  uf = DiffEqDiffTools.UDerivativeWrapper(f,t,p)
  Rosenbrock33ConstantCache(tf,uf,ROS3PConstantCache(real(uBottomEltypeNoUnits),real(tTypeNoUnits)))
end

mutable struct Rosenbrock34Cache{uType,uArrayType,rateType,du2Type,LinuType,vecuType,JType,WType,TabType,TFType,UFType,F,JCType,GCType} <: RosenbrockMutableCache
  u::uType
  uprev::uType
  du::rateType
  du1::rateType
  du2::du2Type
  vectmp::vecuType
  vectmp2::vecuType
  vectmp3::vecuType
  vectmp4::vecuType
  fsalfirst::rateType
  fsallast::rateType
  dT::uArrayType
  J::JType
  W::WType
  tmp::uArrayType
  tab::TabType
  tf::TFType
  uf::UFType
  linsolve_tmp::LinuType
  linsolve_tmp_vec::vecuType
  linsolve::F
  jac_config::JCType
  grad_config::GCType
end

u_cache(c::Rosenbrock34Cache) = (c.dT,c.tmp)
du_cache(c::Rosenbrock34Cache) = (c.k₁,c.k₂,c.k₃,c.du1,c.du2,c.f₁,c.fsalfirst,c.fsallast,c.linsolve_tmp)
jac_cache(c::Rosenbrock34Cache) = (c.J,c.W)
vecu_cache(c::Rosenbrock34Cache) = (c.vectmp,c.vectmp2,c.vectmp3)

function alg_cache(alg::Rodas3,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{true}})
  du = zero(rate_prototype)
  du1 = zero(rate_prototype)
  du2 = zero(rate_prototype)
  vectmp = vec(similar(u,axes(u)))
  vectmp2 = vec(similar(u,axes(u)))
  vectmp3 = vec(similar(u,axes(u)))
  vectmp4 = vec(similar(u,axes(u)))
  fsalfirst = zero(rate_prototype)
  fsallast = zero(rate_prototype)
  dT = similar(u,axes(u))
  if DiffEqBase.has_jac(f) && !DiffEqBase.has_invW(f) && f.jac_prototype != nothing
    W = WOperator(f, dt)
    J = nothing # is J = W.J better?
  else
    J = fill(zero(uEltypeNoUnits),length(u),length(u)) # uEltype?
    W = similar(J)
  end
  tmp = similar(u,axes(u))
  tab = Rodas3ConstantCache(real(uBottomEltypeNoUnits),real(tTypeNoUnits))

  tf = DiffEqDiffTools.TimeGradientWrapper(f,uprev,p)
  uf = DiffEqDiffTools.UJacobianWrapper(f,t,p)
  linsolve_tmp = similar(u,axes(u))
  linsolve_tmp_vec = vec(linsolve_tmp)
  linsolve = alg.linsolve(Val{:init},uf,u)
  grad_config = build_grad_config(alg,f,tf,du1,t)
  jac_config = build_jac_config(alg,f,uf,du1,uprev,u,tmp,du2)
  Rosenbrock34Cache(u,uprev,du,du1,du2,vectmp,vectmp2,vectmp3,vectmp4,
                    fsalfirst,fsallast,dT,J,W,tmp,tab,tf,uf,linsolve_tmp,
                    linsolve_tmp_vec,linsolve,jac_config,grad_config)
end

struct Rosenbrock34ConstantCache{TF,UF,Tab} <: OrdinaryDiffEqConstantCache
  tf::TF
  uf::UF
  tab::Tab
end

function alg_cache(alg::Rodas3,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{false}})
  tf = DiffEqDiffTools.TimeDerivativeWrapper(f,u,p)
  uf = DiffEqDiffTools.UDerivativeWrapper(f,t,p)
  Rosenbrock34ConstantCache(tf,uf,Rodas3ConstantCache(real(uBottomEltypeNoUnits),real(tTypeNoUnits)))
end

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

struct Rosenbrock4ConstantCache{TF,UF,Tab} <: OrdinaryDiffEqConstantCache
  tf::TF
  uf::UF
  tab::Tab
end

mutable struct Rosenbrock4Cache{uType,uArrayType,rateType,du2Type,LinuType,vecuType,JType,WType,TabType,TFType,UFType,F,JCType,GCType} <: RosenbrockMutableCache
  u::uType
  uprev::uType
  du::rateType
  du1::rateType
  du2::du2Type
  vectmp::vecuType
  vectmp2::vecuType
  vectmp3::vecuType
  vectmp4::vecuType
  fsalfirst::rateType
  fsallast::rateType
  dT::uArrayType
  J::JType
  W::WType
  tmp::uArrayType
  tab::TabType
  tf::TFType
  uf::UFType
  linsolve_tmp::LinuType
  linsolve_tmp_vec::vecuType
  linsolve::F
  jac_config::JCType
  grad_config::GCType
end

u_cache(c::Rosenbrock4Cache) = (c.dT,c.tmp)
du_cache(c::Rosenbrock4Cache) = (c.k₁,c.k₂,c.k₃,c.du1,c.du2,c.f₁,c.fsalfirst,c.fsallast,c.linsolve_tmp)
jac_cache(c::Rosenbrock4Cache) = (c.J,c.W)
vecu_cache(c::Rosenbrock4Cache) = (c.vectmp,c.vectmp2,c.vectmp3)

function alg_cache(alg::RosShamp4,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{true}})
  du = zero(rate_prototype)
  du1 = zero(rate_prototype)
  du2 = zero(rate_prototype)
  vectmp = vec(similar(u,axes(u)))
  vectmp2 = vec(similar(u,axes(u)))
  vectmp3 = vec(similar(u,axes(u)))
  vectmp4 = vec(similar(u,axes(u)))
  fsalfirst = zero(rate_prototype)
  fsallast = zero(rate_prototype)
  dT = similar(u,axes(u))
  if DiffEqBase.has_jac(f) && !DiffEqBase.has_invW(f) && f.jac_prototype != nothing
    W = WOperator(f, dt)
    J = nothing # is J = W.J better?
  else
    J = fill(zero(uEltypeNoUnits),length(u),length(u)) # uEltype?
    W = similar(J)
  end
  tmp = similar(u,axes(u))
  tab = RosShamp4ConstantCache(real(uBottomEltypeNoUnits),real(tTypeNoUnits))

  tf = DiffEqDiffTools.TimeGradientWrapper(f,uprev,p)
  uf = DiffEqDiffTools.UJacobianWrapper(f,t,p)
  linsolve_tmp = similar(u,axes(u))
  linsolve_tmp_vec = vec(linsolve_tmp)
  linsolve = alg.linsolve(Val{:init},uf,u)
  grad_config = build_grad_config(alg,f,tf,du1,t)
  jac_config = build_jac_config(alg,f,uf,du1,uprev,u,tmp,du2)
  Rosenbrock4Cache(u,uprev,du,du1,du2,vectmp,vectmp2,vectmp3,vectmp4,
                    fsalfirst,fsallast,dT,J,W,tmp,tab,tf,uf,linsolve_tmp,
                    linsolve_tmp_vec,linsolve,jac_config,grad_config)
end

function alg_cache(alg::RosShamp4,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{false}})
  tf = DiffEqDiffTools.TimeDerivativeWrapper(f,u,p)
  uf = DiffEqDiffTools.UDerivativeWrapper(f,t,p)
  Rosenbrock4ConstantCache(tf,uf,RosShamp4ConstantCache(real(uBottomEltypeNoUnits),real(tTypeNoUnits)))
end

function alg_cache(alg::Veldd4,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{true}})
  du = zero(rate_prototype)
  du1 = zero(rate_prototype)
  du2 = zero(rate_prototype)
  vectmp = vec(similar(u,axes(u)))
  vectmp2 = vec(similar(u,axes(u)))
  vectmp3 = vec(similar(u,axes(u)))
  vectmp4 = vec(similar(u,axes(u)))
  fsalfirst = zero(rate_prototype)
  fsallast = zero(rate_prototype)
  dT = similar(u,axes(u))
  if DiffEqBase.has_jac(f) && !DiffEqBase.has_invW(f) && f.jac_prototype != nothing
    W = WOperator(f, dt)
    J = nothing # is J = W.J better?
  else
    J = fill(zero(uEltypeNoUnits),length(u),length(u)) # uEltype?
    W = similar(J)
  end
  tmp = similar(u,axes(u))
  tab = Veldd4ConstantCache(real(uBottomEltypeNoUnits),real(tTypeNoUnits))

  tf = DiffEqDiffTools.TimeGradientWrapper(f,uprev,p)
  uf = DiffEqDiffTools.UJacobianWrapper(f,t,p)
  linsolve_tmp = similar(u,axes(u))
  linsolve_tmp_vec = vec(linsolve_tmp)
  linsolve = alg.linsolve(Val{:init},uf,u)
  grad_config = build_grad_config(alg,f,tf,du1,t)
  jac_config = build_jac_config(alg,f,uf,du1,uprev,u,tmp,du2)
  Rosenbrock4Cache(u,uprev,du,du1,du2,vectmp,vectmp2,vectmp3,vectmp4,
                    fsalfirst,fsallast,dT,J,W,tmp,tab,tf,uf,linsolve_tmp,
                    linsolve_tmp_vec,linsolve,jac_config,grad_config)
end

function alg_cache(alg::Veldd4,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{false}})
  tf = DiffEqDiffTools.TimeDerivativeWrapper(f,u,p)
  uf = DiffEqDiffTools.UDerivativeWrapper(f,t,p)
  Rosenbrock4ConstantCache(tf,uf,Veldd4ConstantCache(real(uBottomEltypeNoUnits),real(tTypeNoUnits)))
end

function alg_cache(alg::Velds4,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{true}})
  du = zero(rate_prototype)
  du1 = zero(rate_prototype)
  du2 = zero(rate_prototype)
  vectmp = vec(similar(u,axes(u)))
  vectmp2 = vec(similar(u,axes(u)))
  vectmp3 = vec(similar(u,axes(u)))
  vectmp4 = vec(similar(u,axes(u)))
  fsalfirst = zero(rate_prototype)
  fsallast = zero(rate_prototype)
  dT = similar(u,axes(u))
  if DiffEqBase.has_jac(f) && !DiffEqBase.has_invW(f) && f.jac_prototype != nothing
    W = WOperator(f, dt)
    J = nothing # is J = W.J better?
  else
    J = fill(zero(uEltypeNoUnits),length(u),length(u)) # uEltype?
    W = similar(J)
  end
  tmp = similar(u,axes(u))
  tab = Velds4ConstantCache(real(uBottomEltypeNoUnits),real(tTypeNoUnits))

  tf = DiffEqDiffTools.TimeGradientWrapper(f,uprev,p)
  uf = DiffEqDiffTools.UJacobianWrapper(f,t,p)
  linsolve_tmp = similar(u,axes(u))
  linsolve_tmp_vec = vec(linsolve_tmp)
  linsolve = alg.linsolve(Val{:init},uf,u)
  grad_config = build_grad_config(alg,f,tf,du1,t)
  jac_config = build_jac_config(alg,f,uf,du1,uprev,u,tmp,du2)
  Rosenbrock4Cache(u,uprev,du,du1,du2,vectmp,vectmp2,vectmp3,vectmp4,
                    fsalfirst,fsallast,dT,J,W,tmp,tab,tf,uf,linsolve_tmp,
                    linsolve_tmp_vec,linsolve,jac_config,grad_config)
end

function alg_cache(alg::Velds4,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{false}})
  tf = DiffEqDiffTools.TimeDerivativeWrapper(f,u,p)
  uf = DiffEqDiffTools.UDerivativeWrapper(f,t,p)
  Rosenbrock4ConstantCache(tf,uf,Velds4ConstantCache(real(uBottomEltypeNoUnits),real(tTypeNoUnits)))
end

function alg_cache(alg::GRK4T,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{true}})
  du = zero(rate_prototype)
  du1 = zero(rate_prototype)
  du2 = zero(rate_prototype)
  vectmp = vec(similar(u,axes(u)))
  vectmp2 = vec(similar(u,axes(u)))
  vectmp3 = vec(similar(u,axes(u)))
  vectmp4 = vec(similar(u,axes(u)))
  fsalfirst = zero(rate_prototype)
  fsallast = zero(rate_prototype)
  dT = similar(u,axes(u))
  if DiffEqBase.has_jac(f) && !DiffEqBase.has_invW(f) && f.jac_prototype != nothing
    W = WOperator(f, dt)
    J = nothing # is J = W.J better?
  else
    J = fill(zero(uEltypeNoUnits),length(u),length(u)) # uEltype?
    W = similar(J)
  end
  tmp = similar(u,axes(u))
  tab = GRK4TConstantCache(real(uBottomEltypeNoUnits),real(tTypeNoUnits))

  tf = DiffEqDiffTools.TimeGradientWrapper(f,uprev,p)
  uf = DiffEqDiffTools.UJacobianWrapper(f,t,p)
  linsolve_tmp = similar(u,axes(u))
  linsolve_tmp_vec = vec(linsolve_tmp)
  linsolve = alg.linsolve(Val{:init},uf,u)
  grad_config = build_grad_config(alg,f,tf,du1,t)
  jac_config = build_jac_config(alg,f,uf,du1,uprev,u,tmp,du2)
  Rosenbrock4Cache(u,uprev,du,du1,du2,vectmp,vectmp2,vectmp3,vectmp4,
                    fsalfirst,fsallast,dT,J,W,tmp,tab,tf,uf,linsolve_tmp,
                    linsolve_tmp_vec,linsolve,jac_config,grad_config)
end

function alg_cache(alg::GRK4T,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{false}})
  tf = DiffEqDiffTools.TimeDerivativeWrapper(f,u,p)
  uf = DiffEqDiffTools.UDerivativeWrapper(f,t,p)
  Rosenbrock4ConstantCache(tf,uf,GRK4TConstantCache(real(uBottomEltypeNoUnits),real(tTypeNoUnits)))
end

function alg_cache(alg::GRK4A,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{true}})
  du = zero(rate_prototype)
  du1 = zero(rate_prototype)
  du2 = zero(rate_prototype)
  vectmp = vec(similar(u,axes(u)))
  vectmp2 = vec(similar(u,axes(u)))
  vectmp3 = vec(similar(u,axes(u)))
  vectmp4 = vec(similar(u,axes(u)))
  fsalfirst = zero(rate_prototype)
  fsallast = zero(rate_prototype)
  dT = similar(u,axes(u))
  if DiffEqBase.has_jac(f) && !DiffEqBase.has_invW(f) && f.jac_prototype != nothing
    W = WOperator(f, dt)
    J = nothing # is J = W.J better?
  else
    J = fill(zero(uEltypeNoUnits),length(u),length(u)) # uEltype?
    W = similar(J)
  end
  tmp = similar(u,axes(u))
  tab = GRK4AConstantCache(real(uBottomEltypeNoUnits),real(tTypeNoUnits))

  tf = DiffEqDiffTools.TimeGradientWrapper(f,uprev,p)
  uf = DiffEqDiffTools.UJacobianWrapper(f,t,p)
  linsolve_tmp = similar(u,axes(u))
  linsolve_tmp_vec = vec(linsolve_tmp)
  linsolve = alg.linsolve(Val{:init},uf,u)
  grad_config = build_grad_config(alg,f,tf,du1,t)
  jac_config = build_jac_config(alg,f,uf,du1,uprev,u,tmp,du2)
  Rosenbrock4Cache(u,uprev,du,du1,du2,vectmp,vectmp2,vectmp3,vectmp4,
                    fsalfirst,fsallast,dT,J,W,tmp,tab,tf,uf,linsolve_tmp,
                    linsolve_tmp_vec,linsolve,jac_config,grad_config)
end

function alg_cache(alg::GRK4A,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{false}})
  tf = DiffEqDiffTools.TimeDerivativeWrapper(f,u,p)
  uf = DiffEqDiffTools.UDerivativeWrapper(f,t,p)
  Rosenbrock4ConstantCache(tf,uf,GRK4AConstantCache(real(uBottomEltypeNoUnits),real(tTypeNoUnits)))
end

function alg_cache(alg::Ros4LStab,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{true}})
  du = zero(rate_prototype)
  du1 = zero(rate_prototype)
  du2 = zero(rate_prototype)
  vectmp = vec(similar(u,axes(u)))
  vectmp2 = vec(similar(u,axes(u)))
  vectmp3 = vec(similar(u,axes(u)))
  vectmp4 = vec(similar(u,axes(u)))
  fsalfirst = zero(rate_prototype)
  fsallast = zero(rate_prototype)
  dT = similar(u,axes(u))
  if DiffEqBase.has_jac(f) && !DiffEqBase.has_invW(f) && f.jac_prototype != nothing
    W = WOperator(f, dt)
    J = nothing # is J = W.J better?
  else
    J = fill(zero(uEltypeNoUnits),length(u),length(u)) # uEltype?
    W = similar(J)
  end
  tmp = similar(u,axes(u))
  tab = Ros4LStabConstantCache(real(uBottomEltypeNoUnits),real(tTypeNoUnits))

  tf = DiffEqDiffTools.TimeGradientWrapper(f,uprev,p)
  uf = DiffEqDiffTools.UJacobianWrapper(f,t,p)
  linsolve_tmp = similar(u,axes(u))
  linsolve_tmp_vec = vec(linsolve_tmp)
  linsolve = alg.linsolve(Val{:init},uf,u)
  grad_config = build_grad_config(alg,f,tf,du1,t)
  jac_config = build_jac_config(alg,f,uf,du1,uprev,u,tmp,du2)
  Rosenbrock4Cache(u,uprev,du,du1,du2,vectmp,vectmp2,vectmp3,vectmp4,
                    fsalfirst,fsallast,dT,J,W,tmp,tab,tf,uf,linsolve_tmp,
                    linsolve_tmp_vec,linsolve,jac_config,grad_config)
end

function alg_cache(alg::Ros4LStab,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{false}})
  tf = DiffEqDiffTools.TimeDerivativeWrapper(f,u,p)
  uf = DiffEqDiffTools.UDerivativeWrapper(f,t,p)
  Rosenbrock4ConstantCache(tf,uf,Ros4LStabConstantCache(real(uBottomEltypeNoUnits),real(tTypeNoUnits)))
end

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

### Rodas methods

struct Rodas4ConstantCache{TF,UF,Tab} <: OrdinaryDiffEqConstantCache
  tf::TF
  uf::UF
  tab::Tab
end

mutable struct Rodas4Cache{uType,uArrayType,rateType,du2Type,LinuType,vecuType,JType,WType,TabType,TFType,UFType,F,JCType,GCType} <: RosenbrockMutableCache
  u::uType
  uprev::uType
  dense1::rateType
  dense2::rateType
  du::rateType
  du1::rateType
  du2::du2Type
  vectmp::vecuType
  vectmp2::vecuType
  vectmp3::vecuType
  vectmp4::vecuType
  vectmp5::vecuType
  vectmp6::vecuType
  fsalfirst::rateType
  fsallast::rateType
  dT::uArrayType
  J::JType
  W::WType
  tmp::uArrayType
  tab::TabType
  tf::TFType
  uf::UFType
  linsolve_tmp::LinuType
  linsolve_tmp_vec::vecuType
  linsolve::F
  jac_config::JCType
  grad_config::GCType
end

u_cache(c::Rodas4Cache) = (c.dT,c.tmp)
du_cache(c::Rodas4Cache) = (c.k₁,c.k₂,c.k₃,c.du1,c.du2,c.f₁,c.fsalfirst,c.fsallast,c.linsolve_tmp)
jac_cache(c::Rodas4Cache) = (c.J,c.W)
vecu_cache(c::Rodas4Cache) = (c.vectmp,c.vectmp2,c.vectmp3)

function alg_cache(alg::Rodas4,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{true}})
  dense1 = zero(rate_prototype)
  dense2 = zero(rate_prototype)
  du = zero(rate_prototype)
  du1 = zero(rate_prototype)
  du2 = zero(rate_prototype)
  vectmp = vec(similar(u,axes(u)))
  vectmp2 = vec(similar(u,axes(u)))
  vectmp3 = vec(similar(u,axes(u)))
  vectmp4 = vec(similar(u,axes(u)))
  vectmp5 = vec(similar(u,axes(u)))
  vectmp6 = vec(similar(u,axes(u)))
  fsalfirst = zero(rate_prototype)
  fsallast = zero(rate_prototype)
  dT = similar(u,axes(u))
  if DiffEqBase.has_jac(f) && !DiffEqBase.has_invW(f) && f.jac_prototype != nothing
    W = WOperator(f, dt)
    J = nothing # is J = W.J better?
  else
    J = fill(zero(uEltypeNoUnits),length(u),length(u)) # uEltype?
    W = similar(J)
  end
  tmp = similar(u,axes(u))
  tab = Rodas4ConstantCache(real(uBottomEltypeNoUnits),real(tTypeNoUnits))

  tf = DiffEqDiffTools.TimeGradientWrapper(f,uprev,p)
  uf = DiffEqDiffTools.UJacobianWrapper(f,t,p)
  linsolve_tmp = similar(u,axes(u))
  linsolve_tmp_vec = vec(linsolve_tmp)
  linsolve = alg.linsolve(Val{:init},uf,u)
  grad_config = build_grad_config(alg,f,tf,du1,t)
  jac_config = build_jac_config(alg,f,uf,du1,uprev,u,tmp,du2)
  Rodas4Cache(u,uprev,dense1,dense2,du,du1,du2,vectmp,vectmp2,vectmp3,vectmp4,
                    vectmp5,vectmp6,
                    fsalfirst,fsallast,dT,J,W,tmp,tab,tf,uf,linsolve_tmp,
                    linsolve_tmp_vec,linsolve,jac_config,grad_config)
end

function alg_cache(alg::Rodas4,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{false}})
  tf = DiffEqDiffTools.TimeDerivativeWrapper(f,u,p)
  uf = DiffEqDiffTools.UDerivativeWrapper(f,t,p)
  Rodas4ConstantCache(tf,uf,Rodas4ConstantCache(real(uBottomEltypeNoUnits),real(tTypeNoUnits)))
end

function alg_cache(alg::Rodas42,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{true}})
  dense1 = zero(rate_prototype)
  dense2 = zero(rate_prototype)
  du = zero(rate_prototype)
  du1 = zero(rate_prototype)
  du2 = zero(rate_prototype)
  vectmp = vec(similar(u,axes(u)))
  vectmp2 = vec(similar(u,axes(u)))
  vectmp3 = vec(similar(u,axes(u)))
  vectmp4 = vec(similar(u,axes(u)))
  vectmp5 = vec(similar(u,axes(u)))
  vectmp6 = vec(similar(u,axes(u)))
  fsalfirst = zero(rate_prototype)
  fsallast = zero(rate_prototype)
  dT = similar(u,axes(u))
  if DiffEqBase.has_jac(f) && !DiffEqBase.has_invW(f) && f.jac_prototype != nothing
    W = WOperator(f, dt)
    J = nothing # is J = W.J better?
  else
    J = fill(zero(uEltypeNoUnits),length(u),length(u)) # uEltype?
    W = similar(J)
  end
  tmp = similar(u,axes(u))
  tab = Rodas42ConstantCache(real(uBottomEltypeNoUnits),real(tTypeNoUnits))

  tf = DiffEqDiffTools.TimeGradientWrapper(f,uprev,p)
  uf = DiffEqDiffTools.UJacobianWrapper(f,t,p)
  linsolve_tmp = similar(u,axes(u))
  linsolve_tmp_vec = vec(linsolve_tmp)
  linsolve = alg.linsolve(Val{:init},uf,u)
  grad_config = build_grad_config(alg,f,tf,du1,t)
  jac_config = build_jac_config(alg,f,uf,du1,uprev,u,tmp,du2)
  Rodas4Cache(u,uprev,dense1,dense2,du,du1,du2,vectmp,vectmp2,vectmp3,vectmp4,
                    vectmp5,vectmp6,
                    fsalfirst,fsallast,dT,J,W,tmp,tab,tf,uf,linsolve_tmp,
                    linsolve_tmp_vec,linsolve,jac_config,grad_config)
end

function alg_cache(alg::Rodas42,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{false}})
  tf = DiffEqDiffTools.TimeDerivativeWrapper(f,u,p)
  uf = DiffEqDiffTools.UDerivativeWrapper(f,t,p)
  Rodas4ConstantCache(tf,uf,Rodas42ConstantCache(real(uBottomEltypeNoUnits),real(tTypeNoUnits)))
end

function alg_cache(alg::Rodas4P,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{true}})
  dense1 = zero(rate_prototype)
  dense2 = zero(rate_prototype)
  du = zero(rate_prototype)
  du1 = zero(rate_prototype)
  du2 = zero(rate_prototype)
  vectmp = vec(similar(u,axes(u)))
  vectmp2 = vec(similar(u,axes(u)))
  vectmp3 = vec(similar(u,axes(u)))
  vectmp4 = vec(similar(u,axes(u)))
  vectmp5 = vec(similar(u,axes(u)))
  vectmp6 = vec(similar(u,axes(u)))
  fsalfirst = zero(rate_prototype)
  fsallast = zero(rate_prototype)
  dT = similar(u,axes(u))
  if DiffEqBase.has_jac(f) && !DiffEqBase.has_invW(f) && f.jac_prototype != nothing
    W = WOperator(f, dt)
    J = nothing # is J = W.J better?
  else
    J = fill(zero(uEltypeNoUnits),length(u),length(u)) # uEltype?
    W = similar(J)
  end
  tmp = similar(u,axes(u))
  tab = Rodas4PConstantCache(real(uBottomEltypeNoUnits),real(tTypeNoUnits))

  tf = DiffEqDiffTools.TimeGradientWrapper(f,uprev,p)
  uf = DiffEqDiffTools.UJacobianWrapper(f,t,p)
  linsolve_tmp = similar(u,axes(u))
  linsolve_tmp_vec = vec(linsolve_tmp)
  linsolve = alg.linsolve(Val{:init},uf,u)
  grad_config = build_grad_config(alg,f,tf,du1,t)
  jac_config = build_jac_config(alg,f,uf,du1,uprev,u,tmp,du2)
  Rodas4Cache(u,uprev,dense1,dense2,du,du1,du2,vectmp,vectmp2,vectmp3,vectmp4,
                    vectmp5,vectmp6,
                    fsalfirst,fsallast,dT,J,W,tmp,tab,tf,uf,linsolve_tmp,
                    linsolve_tmp_vec,linsolve,jac_config,grad_config)
end

function alg_cache(alg::Rodas4P,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{false}})
  tf = DiffEqDiffTools.TimeDerivativeWrapper(f,u,p)
  uf = DiffEqDiffTools.UDerivativeWrapper(f,t,p)
  Rodas4ConstantCache(tf,uf,Rodas4PConstantCache(real(uBottomEltypeNoUnits),real(tTypeNoUnits)))
end

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

### Rosenbrock5

struct Rosenbrock5ConstantCache{TF,UF,Tab} <: OrdinaryDiffEqConstantCache
  tf::TF
  uf::UF
  tab::Tab
end

mutable struct Rosenbrock5Cache{uType,uArrayType,rateType,du2Type,LinuType,vecuType,JType,WType,TabType,TFType,UFType,F,JCType,GCType} <: RosenbrockMutableCache
  u::uType
  uprev::uType
  dense1::rateType
  dense2::rateType
  du::rateType
  du1::rateType
  du2::du2Type
  vectmp::vecuType
  vectmp2::vecuType
  vectmp3::vecuType
  vectmp4::vecuType
  vectmp5::vecuType
  vectmp6::vecuType
  vectmp7::vecuType
  vectmp8::vecuType
  fsalfirst::rateType
  fsallast::rateType
  dT::uArrayType
  J::JType
  W::WType
  tmp::uArrayType
  tab::TabType
  tf::TFType
  uf::UFType
  linsolve_tmp::LinuType
  linsolve_tmp_vec::vecuType
  linsolve::F
  jac_config::JCType
  grad_config::GCType
end

u_cache(c::Rosenbrock5Cache) = (c.dT,c.tmp)
du_cache(c::Rosenbrock5Cache) = (c.k₁,c.k₂,c.k₃,c.du1,c.du2,c.f₁,c.fsalfirst,c.fsallast,c.linsolve_tmp)
jac_cache(c::Rosenbrock5Cache) = (c.J,c.W)
vecu_cache(c::Rosenbrock5Cache) = (c.vectmp,c.vectmp2,c.vectmp3)

function alg_cache(alg::Rodas5,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{true}})
  dense1 = zero(rate_prototype)
  dense2 = zero(rate_prototype)
  du = zero(rate_prototype)
  du1 = zero(rate_prototype)
  du2 = zero(rate_prototype)
  vectmp = vec(similar(u,axes(u)))
  vectmp2 = vec(similar(u,axes(u)))
  vectmp3 = vec(similar(u,axes(u)))
  vectmp4 = vec(similar(u,axes(u)))
  vectmp5 = vec(similar(u,axes(u)))
  vectmp6 = vec(similar(u,axes(u)))
  vectmp7 = vec(similar(u,axes(u)))
  vectmp8 = vec(similar(u,axes(u)))
  fsalfirst = zero(rate_prototype)
  fsallast = zero(rate_prototype)
  dT = similar(u,axes(u))
  if DiffEqBase.has_jac(f) && !DiffEqBase.has_invW(f) && f.jac_prototype != nothing
    W = WOperator(f, dt)
    J = nothing # is J = W.J better?
  else
    J = fill(zero(uEltypeNoUnits),length(u),length(u)) # uEltype?
    W = similar(J)
  end
  tmp = similar(u,axes(u))
  tab = Rodas5ConstantCache(real(uBottomEltypeNoUnits),real(tTypeNoUnits))

  tf = DiffEqDiffTools.TimeGradientWrapper(f,uprev,p)
  uf = DiffEqDiffTools.UJacobianWrapper(f,t,p)
  linsolve_tmp = similar(u,axes(u))
  linsolve_tmp_vec = vec(linsolve_tmp)
  linsolve = alg.linsolve(Val{:init},uf,u)
  grad_config = build_grad_config(alg,f,tf,du1,t)
  jac_config = build_jac_config(alg,f,uf,du1,uprev,u,tmp,du2)
  Rosenbrock5Cache(u,uprev,dense1,dense2,du,du1,du2,vectmp,vectmp2,vectmp3,vectmp4,
                    vectmp5,vectmp6,vectmp7,vectmp8,
                    fsalfirst,fsallast,dT,J,W,tmp,tab,tf,uf,linsolve_tmp,
                    linsolve_tmp_vec,linsolve,jac_config,grad_config)
end

function alg_cache(alg::Rodas5,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{false}})
  tf = DiffEqDiffTools.TimeDerivativeWrapper(f,u,p)
  uf = DiffEqDiffTools.UDerivativeWrapper(f,t,p)
  Rosenbrock5ConstantCache(tf,uf,Rodas5ConstantCache(real(uBottomEltypeNoUnits),real(tTypeNoUnits)))
end


@testset "smc.jl" begin
    @testset "SMC constructor" begin
        sampler = AdvancedPS.SMC(10)
        @test sampler.nparticles == 10
        @test sampler.resampler == AdvancedPS.ResampleWithESSThreshold()

        sampler = AdvancedPS.SMC(15, 0.6)
        @test sampler.nparticles == 15
        @test sampler.resampler ===
              AdvancedPS.ResampleWithESSThreshold(AdvancedPS.resample_systematic, 0.6)

        sampler = AdvancedPS.SMC(20, AdvancedPS.resample_multinomial, 0.6)
        @test sampler.nparticles == 20
        @test sampler.resampler ===
              AdvancedPS.ResampleWithESSThreshold(AdvancedPS.resample_multinomial, 0.6)

        sampler = AdvancedPS.SMC(25, AdvancedPS.resample_systematic)
        @test sampler.nparticles == 25
        @test sampler.resampler === AdvancedPS.resample_systematic
    end

    # Smoke tests
    @testset "models" begin
        mutable struct NormalModel <: AbstractMCMC.AbstractModel
            a::Float64
            b::Float64

            NormalModel() = new()
        end

        function (m::NormalModel)(rng::Random.AbstractRNG)
            # First latent variable.
            m.a = a = rand(rng, Normal(4, 5))

            # First observation.
            AdvancedPS.observe(Normal(a, 2), 3)

            # Second latent variable.
            m.b = b = rand(rng, Normal(a, 1))

            # Second observation.
            return AdvancedPS.observe(Normal(b, 2), 1.5)
        end

        sample(NormalModel(), AdvancedPS.SMC(100))

        # failing test
        mutable struct FailSMCModel <: AbstractMCMC.AbstractModel
            a::Float64
            b::Float64

            FailSMCModel() = new()
        end

        function (m::FailSMCModel)(rng::Random.AbstractRNG)
            m.a = a = rand(rng, Normal(4, 5))
            m.b = b = rand(rng, Normal(a, 1))
            if a >= 4
                AdvancedPS.observe(Normal(b, 2), 1.5)
            end
        end

        @test_throws ErrorException sample(FailSMCModel(), AdvancedPS.SMC(100))
    end

    @testset "logevidence" begin
        Random.seed!(100)

        mutable struct TestModel <: AbstractMCMC.AbstractModel
            a::Float64
            x::Bool
            b::Float64
            c::Float64

            TestModel() = new()
        end

        function (m::TestModel)(rng::Random.AbstractRNG)
            # First hidden variables.
            m.a = rand(rng, Normal(0, 1))
            m.x = x = rand(rng, Bernoulli(1))
            m.b = rand(rng, Gamma(2, 3))

            # First observation.
            AdvancedPS.observe(Bernoulli(x / 2), 1)

            # Second hidden variable.
            m.c = rand(rng, Beta())

            # Second observation.
            return AdvancedPS.observe(Bernoulli(x / 2), 0)
        end

        chains_smc = sample(TestModel(), AdvancedPS.SMC(100))

        @test all(isone(p.f.x) for p in chains_smc.trajectories)
        @test chains_smc.logevidence ≈ -2 * log(2)
    end

    @testset "PG constructor" begin
        sampler = AdvancedPS.PG(10)
        @test sampler.nparticles == 10
        @test sampler.resampler == AdvancedPS.ResampleWithESSThreshold()

        sampler = AdvancedPS.PG(60, 0.6)
        @test sampler.nparticles == 60
        @test sampler.resampler ===
              AdvancedPS.ResampleWithESSThreshold(AdvancedPS.resample_systematic, 0.6)

        sampler = AdvancedPS.PG(80, AdvancedPS.resample_multinomial, 0.6)
        @test sampler.nparticles == 80
        @test sampler.resampler ===
              AdvancedPS.ResampleWithESSThreshold(AdvancedPS.resample_multinomial, 0.6)

        sampler = AdvancedPS.PG(100, AdvancedPS.resample_systematic)
        @test sampler.nparticles == 100
        @test sampler.resampler === AdvancedPS.resample_systematic
    end

    @testset "logevidence" begin
        Random.seed!(100)

        mutable struct TestModel <: AbstractMCMC.AbstractModel
            a::Float64
            x::Bool
            b::Float64
            c::Float64

            TestModel() = new()
        end

        function (m::TestModel)(rng::Random.AbstractRNG)
            # First hidden variables.
            m.a = rand(rng, Normal(0, 1))
            m.x = x = rand(rng, Bernoulli(1))
            m.b = rand(rng, Gamma(2, 3))

            # First observation.
            AdvancedPS.observe(Bernoulli(x / 2), 1)

            # Second hidden variable.
            m.c = rand(rng, Beta())

            # Second observation.
            return AdvancedPS.observe(Bernoulli(x / 2), 0)
        end

        chains_pg = sample(TestModel(), AdvancedPS.PG(10), 100)

        @test all(isone(p.trajectory.f.x) for p in chains_pg)
        @test mean(x.logevidence for x in chains_pg) ≈ -2 * log(2) atol = 0.01
    end

    @testset "Replay reference" begin
        mutable struct Model <: AbstractMCMC.AbstractModel
            a::Float64
            b::Float64

            Model() = new()
        end

        function (m::Model)(rng)
            m.a = rand(rng, Normal())
            AdvancedPS.observe(Normal(), m.a)

            m.b = rand(rng, Normal())
            return AdvancedPS.observe(Normal(), m.b)
        end

        pg = AdvancedPS.PG(1)
        first, second = sample(Model(), pg, 2)

        first_model = first.trajectory.f
        second_model = second.trajectory.f

        # Single Particle - must be replaying
        @test first_model.a ≈ second_model.a
        @test first_model.b ≈ second_model.b
        @test first.logevidence ≈ second.logevidence
    end
end

# @testset "pmmh.jl" begin
#     @turing_testset "pmmh constructor" begin
#         N = 2000
#         s1 = PMMH(N, SMC(10, :s), MH(1,(:m, s -> Normal(s, sqrt(1)))))
#         s2 = PMMH(N, SMC(10, :s), MH(1, :m))
#         s3 = PIMH(N, SMC())
#
#         c1 = sample(gdemo_default, s1)
#         c2 = sample(gdemo_default, s2)
#         c3 = sample(gdemo_default, s3)
#     end
#     @numerical_testset "pmmh inference" begin
#         alg = PMMH(2000, SMC(20, :m), MH(1, (:s, GKernel(1))))
#         chain = sample(gdemo_default, alg)
#         check_gdemo(chain, atol = 0.1)
#
#         # PMMH with prior as proposal
#         alg = PMMH(2000, SMC(20, :m), MH(1, :s))
#         chain = sample(gdemo_default, alg)
#         check_gdemo(chain, atol = 0.1)
#
#         # PIMH
#         alg = PIMH(2000, SMC())
#         chain = sample(gdemo_default, alg)
#         check_gdemo(chain)
#
#         # MoGtest
#         pmmh = PMMH(2000,
#             SMC(10, :z1, :z2, :z3, :z4),
#             MH(1, :mu1, :mu2))
#         chain = sample(MoGtest_default, pmmh)
#
#         check_MoGtest_default(chain, atol = 0.1)
#     end
# end

# @testset "ipmcmc.jl" begin
#     @turing_testset "ipmcmc constructor" begin
#         Random.seed!(125)
#
#         N = 50
#         s1 = IPMCMC(10, N, 4, 2)
#         s2 = IPMCMC(10, N, 4)
#
#         c1 = sample(gdemo_default, s1)
#         c2 = sample(gdemo_default, s2)
#     end
#     @numerical_testset "ipmcmc inference" begin
#         alg = IPMCMC(30, 500, 4)
#         chain = sample(gdemo_default, alg)
#         check_gdemo(chain)
#
#         alg2 = IPMCMC(15, 100, 10)
#         chain2 = sample(MoGtest_default, alg2)
#         check_MoGtest_default(chain2, atol = 0.2)
#     end
# end


function main!(r::Robot)
    t = 0
    k = 0
    a = 0
    while !isborder(r,Ost)
        while !isborder(r, Nord)
            move!(r,Nord)
            if ismarker(r)
                k = k + 1
                t = t + temperature(r)
            end
        end
        goback!(r)
        move!(r, Ost)
        if ismarker(r)
            k = k + 1
            t = t + temperature(r)
        end
    end
    while !isborder(r, Nord)
        move!(r, Nord)
        if ismarker(r)
            k = k + 1
            t = t + temperature(r)
        end
    end
    if k > 0
        a = t/k
        println(a)
    end
end


# draw histogram of sensor data 600 as time series

module DrawTimeSeries600
    using Plots, DataFrames, CSV
    pyplot()

    function main(is_test=false)
        data_path = joinpath(split(@__FILE__, "src")[1], "data/sensor_data_600.txt")
        df_600_mm = CSV.read(data_path, DataFrame,
                             header=["date", "time", "ir", "lidar"],
                             delim=' ')

        time_idx = Array(0:length(df_600_mm.lidar)-1)

        plot(time_idx, df_600_mm.lidar, lebel="time series")

        if is_test == false
            save_path = joinpath(split(@__FILE__, "src")[1], "src/prob_stats/complex_dist/time_series_600mm/time_series_600_mm.png")
            savefig(save_path)
        end
    end
end


import Mercury as Hg
import Mercury: LoopRateLimiter
using LinearAlgebra
using StaticArrays
using Statistics
using Test

# @testset "Rate Limited Loops" begin
function runrate(f, lrl::LoopRateLimiter, time_s = 1)
    N = round(Int, lrl.rate * time_s)
    times = zeros(N)
    Hg.reset!(lrl)
    for i = 1:N
        times[i] = @elapsed begin
            Hg.startloop(lrl)
            f()
            sleep(lrl)
        end
    end
    return 1 ./ times
end

@generated function runrate_macro1(f, ::Val{rate}, time_s = 1) where {rate}
    quote
        N = round(Int, $rate * time_s)
        lrl = LoopRateLimiter(rate)
        t_start = time_ns()
        Hg.@rate for i = 1:N
            f()
        end $rate
        t_elapsed = (time_ns() - t_start) * 1e-9
        return N / t_elapsed
    end
end

function runrate_macro2(f, rate, time_s = 1)
    N = round(Int, rate * time_s)
    lrl = LoopRateLimiter(rate)
    t_start = time_ns()
    Hg.@rate for i = 1:N
        f()
    end lrl
    t_elapsed = (time_ns() - t_start) * 1e-9
    return N / t_elapsed
end

@generated function runrate_macro3(f, ::Val{rate}, time_s = 1) where {rate}
    quote
        N = round(Int, $rate * time_s)
        lrl = LoopRateLimiter(rate)
        t_start = time_ns()
        i = 0
        Hg.@rate while (i < N)
            f()
            i += 1
        end $rate
        t_elapsed = (time_ns() - t_start) * 1e-9
        return N / t_elapsed
    end
end

function runrate_macro4(f, rate, time_s = 1)
    N = round(Int, rate * time_s)
    lrl = LoopRateLimiter(rate)
    t_start = time_ns()
    i = 0
    Hg.@rate while (i < N)
        f()
        i += 1
    end lrl
    t_elapsed = (time_ns() - t_start) * 1e-9
    return N / t_elapsed
end

##
function mykernel()
    A = @SMatrix randn(10, 10)
    B = @SMatrix randn(10, 10)
    C = A'A + B'B
    x = @SVector randn(10)
    cholA = cholesky(Symmetric(C))
    C \ x
end
1 / @elapsed mykernel()

@testset "Rate limiter tests" begin
    rate = 1000
    test_time = 0.1 # sec
    lrl = LoopRateLimiter(rate)
    @test Hg.doesblock(lrl)
    error = (rate - median(runrate(mykernel, lrl, test_time))) / rate * 100
    @test error < 10  # less than 5% error

    error = (rate - runrate_macro1(mykernel, Val(rate), test_time)) / rate * 100
    @test error < 10  # less than 5% error
    error = (rate - runrate_macro2(mykernel, rate, test_time)) / rate * 100
    @test error < 10  # less than 5% error
    error = (rate - runrate_macro3(mykernel, Val(rate), test_time)) / rate * 100
    @test error < 10  # less than 5% error
    error = (rate - runrate_macro4(mykernel, rate, test_time)) / rate * 100
    @test error < 10  # less than 5% error

    rate = 40
    test_time = 0.1 # sec
    thresh = 6  # percent error
    lrl = LoopRateLimiter(rate)
    @test !Hg.doesblock(lrl)
    runrate(mykernel, lrl, test_time)
    error = (rate - median(runrate(mykernel, lrl, test_time))) / rate * 100
    @test error < thresh

    error = (rate - runrate_macro1(mykernel, Val(rate), test_time)) / rate * 100
    @test error < thresh
    error = (rate - runrate_macro2(mykernel, rate, test_time)) / rate * 100
    @test error < thresh
    error = (rate - runrate_macro3(mykernel, Val(rate), test_time)) / rate * 100
    @test error < thresh
    error = (rate - runrate_macro4(mykernel, rate, test_time)) / rate * 100
    @test error < thresh

end


using Random

sym(x) = Symbol(x)

function populate_dependency!(doc, comp)
    empty!(comp.deps)
    for x in list_encapsulation(doc, comp)
        push!(comp.deps, sym(x["component"]))
    end
end

function add_component!(doc, name, node, populate=true)
    comp = Component(name, node, Set{Symbol}())
    populate && populate_dependency!(doc, comp)
    push!(doc.comps, comp)
    comp
end

function add_connection!(doc, c1, c2, vars)
    conn = Connection(c1, c2, vars)
    push!(doc.conns, conn)
    conn
end

"""
    load_cellml loads a CellML file and populates a Document structures with
    the components and connections.

    If resolve = true, load_cellml also resolve the imported components.
"""
function load_cellml(path; resolve=true)
    xml = readxml(path)
    doc = Document(path, [xml], Component[], Connection[])

    for c in list_components(xml)
        add_component!(doc, sym(c["name"]), c)
    end

    for conn in list_connections(xml)
        c1, c2 = sym.(components_of(get_connection_component(conn)))
        vars = [make_var(c1,v1) => make_var(c2,v2)
                for (v1,v2) in variables_of.(list_connection_variables(conn))]
        add_connection!(doc, c1, c2, vars)
    end

    resolve && resolve_imports!(doc)

    return doc
end

"""
    find_component finds a component by name
"""
function find_component(comps::Array{Component}, name)
    i = findfirst(x -> x.name == Symbol(name), comps)
    return i == nothing ? nothing : comps[i]
end

find_component(doc::Document, name) = find_component(doc.comps, name)

# implicit_name(sym) = Symbol("$(string(sym))_$(randstring(6))")
implicit_name(sym) = Symbol("$(string(sym))")

"""
    resolve_imports recursivelly resolves the imported components of doc.
"""
function resolve_imports!(doc::Document)
    for ϵ in list_imports(doc)  # εισαγωγή == import
        href = ϵ["xlink:href"]
        @info "importing $href"
        path = joinpath(splitdir(doc.path)[1], href)
        child = load_cellml(path)
        append!(doc.xmls, child.xmls)

        L = Dict{Symbol, Symbol}()      # L keeps the import names of components
        D = Set{Symbol}()               # D is the set of imported dependencies
        C = Set{Symbol}()

        # first, import explicitly imported component
        for σ in list_import_components(ϵ)  # συστατικό == component
            old_name = sym(σ["component_ref"])
            new_name = sym(σ["name"])
            old_comp = find_component(child, old_name)

            add_component!(doc, new_name, old_comp.node)

            D = D ∪ old_comp.deps
            L[old_name] = new_name
            push!(C, old_name)
        end

        # second, import implicitly imported component, as defined by closure
        #for c in find_closure(child, collect(keys(L)))
        for old_name in D
            if old_name ∉ C
                old_comp = find_component(child, old_name)
                new_name = implicit_name(old_name) # generates a name with a random suffix to prevent name collision
                add_component!(doc, new_name, old_comp.node, false)
                L[old_name] = new_name
            end
        end

        # third, we also need to import connections from the child CellML file
        for conn in connections(child)
            c1, c2 = components(conn)
            if haskey(L,c1) && haskey(L,c2)
                d1, d2 = L[c1], L[c2]
                vars = [Var(d1,v1.var) => Var(d2,v2.var) for (v1,v2) in conn.vars]
                add_connection!(doc, d1, d2, vars)
            end
        end
    end
end

"""
    find_closure finds the transitive closure of a list of componenets (l) minus
    the list itself, i.e., it returns the list of components in doc which are
    reachable through a chain of connections starting from any component in l.
"""
function find_closure(doc::Document, l)
    n = copy(l)

    done = false
    while !done
        done = true
        for (c1,c2) in components.(connections(doc))
            if c1 ∈ n && c2 ∉ n
                push!(n, c2)
                done = false
            elseif c1 ∉ n && c2 ∈ n
                push!(n, c1)
                done = false
            end
        end
    end

    [x for x in n if x ∉ l]
end

"""
Given a path to a directory containing multiple CellML files.

Returns the top level files.

Used as a helper function with the CellML Physiome Model repositories.
"""
function list_top_cellml_files(dir)
    files = filter(f->endswith(f, ".cellml"), readdir(dir))

    if length(files) == 1
        return joinpath.(dir, files)
    end

    imported = Set{String}()
    for f in files
        xml = readxml(joinpath(dir, f))
        for n in CellMLToolkit.list_imports(xml)
            push!(imported, n["xlink:href"])
        end
    end

    if !isempty(imported)
        printstyled("imported files are $imported\n"; color=:yellow)
    end

    joinpath.(dir, [f for f in files if f ∉ imported])
end


struct ZygoteDiffBackend <: AbstractDiffBackend end

function Manifolds._gradient(f, p, ::ZygoteDiffBackend)
    return Zygote.gradient(f, p)[1]
end

function Manifolds._gradient!(f, X, p, ::ZygoteDiffBackend)
    return copyto!(X, Zygote.gradient(f, p)[1])
end

if default_differential_backend() === NoneDiffBackend()
    set_default_differential_backend!(ZygoteDiffBackend())
end


struct Nested{T}
    o::T
end

Base.iterate(u::Nested, i = u.o) = i isa AbstractObservable ? (i, i[]) : nothing

mutable struct Flatten <: AbstractObservable{Any}
    o::AbstractObservable
    list::Vector{Tuple{AbstractObservable, Function}}
    inner_id::String
    output::Observable{Any}
    pair::ObservablePair
    function Flatten(o)
        inner_obs = foldl((x,y)->y, Nested(o))
        inner_id = obsid(inner_obs)
        output = Observable{Any}(inner_obs[])
        p = ObservablePair(inner_obs, output)
        n = new(o, Tuple{AbstractObservable, Function}[], inner_id, output, p)
        for i in Nested(o)
            f = updater(i, n)
            on(f, i)
            push!(n.list, (i, f))
        end
        n
    end
end

flatten(x) = x
flatten(obs::AbstractObservable) = observe(Flatten(obs))

observe(o::Flatten) = o.output

function updater(i, f::Flatten)
    function (val)
        ind = findfirst(t -> obsid(t[1]) == obsid(i), f.list)
        if ind !== nothing
            for ii in length(f.list):-1:(ind+1)
                obs, func = pop!(f.list)
                off(obs, func)
            end
            for outer i in Nested(val)
                func = updater(i, f)
                on(func, i)
                push!(f.list, (i, func))
            end
            if obsid(i) != f.inner_id
                off(f.pair)
                f[] = i[]
                f.pair = ObservablePair(i, f.output)
                f.inner_id = obsid(i)
            end
        end
    end
end


data = split(read(stdin, String), "\n\n", keepempty=false)

mutable struct Tile
    name
    l
    u
    r
    d
    rotation
    borders
end

new_tile(name, l, u, r, d; borders=[]) = Tile(name, l, u, r, d, 0, borders)

# Reverse rotation (ltr)
tiles = []
for tile in data
    tile = split(tile, '\n', keepempty=false)
    name = parse(Int, tile[1][6:end - 1])

    tile = [tile[i][j] for i = 1:length(tile), j = 1:length(tile[1])]
    l = tile[end:-1:2, 1]
    u = tile[2, :]
    r = tile[2:end, end]
    d = tile[end, end:-1:begin]

    push!(tiles, new_tile(name, l, u, r, d, borders=[l, u, r, d]))
end

# Find borders
count = 0
# ids[tileid] = #edge
ids = Dict()
for tile in tiles
    global count

    borders = [tile.l, tile.u, tile.r, tile.d]
    for (borderi, border) in enumerate(borders)
        match = false
        for tile2 in tiles
            if tile2 == tile
                continue
            end

            borders2 = [tile2.l, tile2.u, tile2.r, tile2.d]

            for border2 in borders2
                # Found matching border
                if border2 == border || border2 == border[end:-1:begin]
                    match = true
                    break
                end
            end

            if match
                break
            end
        end

        if !match
            count += 1

            if haskey(ids, tile.name)
                push!(ids[tile.name], borderi)
            else
                ids[tile.name] = [borderi]
            end
        end
    end
end

corners = [tile for (tile, v) in ids if length(v) == 2]

@assert length(corners) == 4

prod = reduce(*, corners)

println("Part 1 : $prod")

# TODO : Remove
# # Top left corner
# first = [t for t in tiles if t.name == corners[end]][begin]
# println("First : $(first.name)")

# current = first
# for tile in tiles
#     global current

#     if tile == current
#         continue
#     end

#     match = false
#     for (borderi, border) in enumerate(tile.borders)
#         if border == current.borders[3]
#             # Not flipped
#             println(tile.name)
#             match = true
#             break
#         elseif border == current.borders[3][end:-1:begin]
#             # Flipped
#             println("$(tile.name) flipped $borderi")
#             match = true
#             break
#         end
#     end

#     if match
#         break
#     end
# end

# TODO : Remove
# Old part 1
# visited = repeat([false], length(tiles))

# # map[i, j] = tile at i, j
# map = Dict()
# map[(0, 0)] = tiles[begin]
# visited[begin] = true

# println(tiles[begin])
# println(tiles[end])

# mini = -0xffff
# minj = -0xffff
# maxi = 0xffff
# maxj = 0xffff

# # Left
# # rot_id : left = 1, down = 2, ...
# function solve(i, j, ij_new, rot_id)
#     global map, visited, tiles
#     global mini, minj, maxi, maxj

#     if !(mini <= i <= maxi) || !(minj <= j <= maxj)
#         return false
#     end

#     current = map[(i, j)]
#     # if current == nothing
#     #     return false
#     # end

#     cedges = [current.l, current.u, current.r, current.d]
#     if !haskey(map, ij_new)
#         for (tilei, (visit, tile)) in enumerate(zip(visited, tiles))
#             if !visit
#                 matches = false
#                 newedges = [tile.l, tile.u, tile.r, tile.d]

#                 # Test each rotation possibility
#                 for rot = 0:3
#                     if newedges[1 + rot] == cedges[1 + (rot_id + current.rotation - 1) % 4][end:-1:begin] || newedges[1 + rot] == cedges[1 + (rot_id + current.rotation - 1) % 4]
#                         matches = true
#                         tile.rotation = rot
#                         break
#                     end
#                 end

#                 if matches
#                     visited[tilei] = true
#                     map[ij_new] = tile
#                     return true
#                 end
#             end
#         end
#     end

#     # map[ij_new] = nothing

#     if rot_id == 1
#         minj = max(minj, j)
#     elseif rot_id == 2
#         mini = max(mini, i)
#     elseif rot_id == 3
#         maxj = min(maxj, j)
#     elseif rot_id == 4
#         maxi = min(maxi, i)
#     end

#     return false
# end

# function solveall(i, j)
#     solve(i, j, (i, j - 1), 1) && solveall(i, j - 1)
#     solve(i, j, (i - 1, j), 2) && solveall(i - 1, j)
#     solve(i, j, (i, j + 1), 3) && solveall(i, j + 1)
#     solve(i, j, (i + 1, j), 4) && solveall(i + 1, j)
# end

# # function solvefrom(i, j)
# #     global map, visited, tiles

# #     solveall(i, j)
# # end

# # solve(0, 0, (-1, 0), 2)
# # solve(0, 0, (-1, 0), 2)
# # solvefrom(0, 0, (-1, 0), 2)
# solveall(0, 0)

# # last = (0, 0)
# # for i = 1:dim(tiles, 1)
# #     for j = 1:dim(tiles, 2)
# #         solveall(last)
# #     end
# # end

# @assert all(visited) "Error : All tiles not visited"

# println()
# println(map)

# # for (k, v) in map
# #     if v != nothing
# #         (i, j) = k

# #         global mini = min(mini, i)
# #         global minj = min(minj, j)
# #         global maxi = max(maxi, i)
# #         global maxj = max(maxj, j)
# #     end
# # end

# println(mini)
# println(minj)
# println(maxi)
# println(maxj)

# # println("> $(map[(1, 0)])")

# prod = map[(mini, minj)] * map[(maxi, minj)] * map[(maxi, maxj)] * map[(mini, maxj)]
# println("Part 1 : $prod")