## Integrator Dispatches

# Can get rid of an allocation here with a function
# get_tmp_arr(integrator.cache) which gives a pointer to some
# cache array which can be modified.

@inline function _searchsortedfirst(v::AbstractVector, x, lo::Integer, forward::Bool)
    u = oftype(lo, 1)
    lo = lo - u
    hi = length(v) + u
    @inbounds while lo < hi - u
        m = (lo + hi) >>> 1
        @inbounds if (forward && v[m] < x) || (!forward && v[m] > x)
            lo = m
        else
            hi = m
        end
    end
    return hi
end

@inline function _searchsortedlast(v::AbstractVector, x, lo::Integer, forward::Bool)
  u = oftype(lo, 1)
  lo = lo - u
  hi = length(v) + u
  @inbounds while lo < hi - u
      m = (lo + hi) >>> 1
      @inbounds if (forward && v[m] > x) || (!forward && v[m] < x)
          hi = m
      else
          lo = m
      end
  end
  return lo
end

@inline function _ode_addsteps!(integrator,f=integrator.f,always_calc_begin = false,allow_calc_end = true,force_calc_end = false)
  if !(typeof(integrator.cache) <: CompositeCache)
    DiffEqBase.addsteps!(integrator.k,integrator.tprev,integrator.uprev,integrator.u,
                  integrator.dt,f,integrator.p,integrator.cache,
                  always_calc_begin,allow_calc_end,force_calc_end)
  else
    DiffEqBase.addsteps!(integrator.k,integrator.tprev,integrator.uprev,integrator.u,
                  integrator.dt,f,integrator.p,
                  @inbounds(integrator.cache.caches[integrator.cache.current]),
                  always_calc_begin,allow_calc_end,force_calc_end)
  end
end
@inline DiffEqBase.addsteps!(integrator::ODEIntegrator,args...) = _ode_addsteps!(integrator,args...)

@inline function ode_interpolant(Θ,integrator::DiffEqBase.DEIntegrator,idxs,deriv)
  DiffEqBase.addsteps!(integrator)
  if !(typeof(integrator.cache) <: CompositeCache)
    val = ode_interpolant(Θ,integrator.dt,integrator.uprev,integrator.u,integrator.k,integrator.cache,idxs,deriv)
  else
    val = composite_ode_interpolant(Θ,integrator,integrator.cache.caches,integrator.cache.current,idxs,deriv)
  end
  val
end

@generated function composite_ode_interpolant(Θ, integrator, caches::T, current, idxs, deriv) where {T <: Tuple}
  expr = Expr(:block)
  for i in 1:length(T.types)
    push!(expr.args, quote
      if $i == current
        return ode_interpolant(Θ,integrator.dt,integrator.uprev,integrator.u,integrator.k,caches[$i],idxs,deriv)
      end
    end)
  end
  push!(expr.args, quote
    throw("Cache $current is not available. There are only $(length(caches)) caches.")
  end)
  return expr
end

@inline function ode_interpolant!(val,Θ,integrator::DiffEqBase.DEIntegrator,idxs,deriv)
  DiffEqBase.addsteps!(integrator)
  if !(typeof(integrator.cache) <: CompositeCache)
    ode_interpolant!(val,Θ,integrator.dt,integrator.uprev,integrator.u,integrator.k,integrator.cache,idxs,deriv)
  else
    ode_interpolant!(val,Θ,integrator.dt,integrator.uprev,integrator.u,integrator.k,integrator.cache.caches[integrator.cache.current],idxs,deriv)
  end
end

@generated function composite_ode_interpolant!(val, Θ, integrator, caches::T, current, idxs, deriv) where {T <: Tuple}
  expr = Expr(:block)
  for i in 1:length(T.types)
    push!(expr.args, quote
      if $i == current
        return ode_interpolant!(val,Θ,integrator.dt,integrator.uprev,integrator.u,integrator.k,caches[$i],idxs,deriv)
      end
    end)
  end
  push!(expr.args, quote
    throw("Cache $current is not available. There are only $(length(caches)) caches.")
  end)
  return expr
end

@inline function current_interpolant(t::Number,integrator::DiffEqBase.DEIntegrator,idxs,deriv)
  Θ = (t-integrator.tprev)/integrator.dt
  ode_interpolant(Θ,integrator,idxs,deriv)
end

@inline function current_interpolant(t,integrator::DiffEqBase.DEIntegrator,idxs,deriv)
  Θ = (t.-integrator.tprev)./integrator.dt
  [ode_interpolant(ϕ,integrator,idxs,deriv) for ϕ in Θ]
end

@inline function current_interpolant!(val,t::Number,integrator::DiffEqBase.DEIntegrator,idxs,deriv)
  Θ = (t-integrator.tprev)/integrator.dt
  ode_interpolant!(val,Θ,integrator,idxs,deriv)
end

@inline function current_interpolant!(val,t,integrator::DiffEqBase.DEIntegrator,idxs,deriv)
  Θ = (t.-integrator.tprev)./integrator.dt
  [ode_interpolant!(val,ϕ,integrator,idxs,deriv) for ϕ in Θ]
end

@inline function current_extrapolant(t::Number,integrator::DiffEqBase.DEIntegrator,idxs=nothing,deriv=Val{0})
  Θ = (t-integrator.tprev)/(integrator.t-integrator.tprev)
  ode_extrapolant(Θ,integrator,idxs,deriv)
end

@inline function current_extrapolant!(val,t::Number,integrator::DiffEqBase.DEIntegrator,idxs=nothing,deriv=Val{0})
  Θ = (t-integrator.tprev)/(integrator.t-integrator.tprev)
  ode_extrapolant!(val,Θ,integrator,idxs,deriv)
end

@inline function current_extrapolant(t::AbstractArray,integrator::DiffEqBase.DEIntegrator,idxs=nothing,deriv=Val{0})
  Θ = (t.-integrator.tprev)./(integrator.t-integrator.tprev)
  [ode_extrapolant(ϕ,integrator,idxs,deriv) for ϕ in Θ]
end

@inline function current_extrapolant!(val,t,integrator::DiffEqBase.DEIntegrator,idxs=nothing,deriv=Val{0})
  Θ = (t.-integrator.tprev)./(integrator.t-integrator.tprev)
  [ode_extrapolant!(val,ϕ,integrator,idxs,deriv) for ϕ in Θ]
end

@inline function ode_extrapolant!(val,Θ,integrator::DiffEqBase.DEIntegrator,idxs,deriv)
  DiffEqBase.addsteps!(integrator)
  if !(typeof(integrator.cache) <: CompositeCache)
    ode_interpolant!(val,Θ,integrator.t-integrator.tprev,integrator.uprev2,integrator.uprev,integrator.k,integrator.cache,idxs,deriv)
  else
    composite_ode_extrapolant!(val,Θ,integrator,integrator.cache.caches,integrator.cache.current,idxs,deriv)
  end
end

@generated function composite_ode_extrapolant!(val, Θ, integrator, caches::T, current, idxs, deriv) where {T <: Tuple}
  expr = Expr(:block)
  for i in 1:length(T.types)
    push!(expr.args, quote
      if $i == current
        return ode_interpolant!(val,Θ,integrator.t-integrator.tprev,integrator.uprev2,integrator.uprev,integrator.k,caches[$i],idxs,deriv)
      end
    end)
  end
  push!(expr.args, quote
    throw("Cache $current is not available. There are only $(length(caches)) caches.")
  end)
  return expr
end

@inline function ode_extrapolant(Θ,integrator::DiffEqBase.DEIntegrator,idxs,deriv)
  DiffEqBase.addsteps!(integrator)
  if !(typeof(integrator.cache) <: CompositeCache)
    ode_interpolant(Θ,integrator.t-integrator.tprev,integrator.uprev2,integrator.uprev,integrator.k,integrator.cache,idxs,deriv)
  else
    composite_ode_extrapolant(Θ,integrator,integrator.cache.caches,integrator.cache.current,idxs,deriv)
  end
end

@generated function composite_ode_extrapolant(Θ, integrator, caches::T, current, idxs, deriv) where {T <: Tuple}
  expr = Expr(:block)
  for i in 1:length(T.types)
    push!(expr.args, quote
      if $i == current
        return ode_interpolant(Θ,integrator.t-integrator.tprev,integrator.uprev2,integrator.uprev,integrator.k,caches[$i],idxs,deriv)
      end
    end)
  end
  push!(expr.args, quote
    throw("Cache $current is not available. There are only $(length(caches)) caches.")
  end)
  return expr
end

"""
ode_interpolation(tvals,ts,timeseries,ks)

Get the value at tvals where the solution is known at the
times ts (sorted), with values timeseries and derivatives ks
"""
function ode_interpolation(tvals,id,idxs,deriv,p,continuity::Symbol=:left)
  @unpack ts,timeseries,ks,f,cache = id
  @inbounds tdir = sign(ts[end]-ts[1])
  idx = sortperm(tvals,rev=tdir<0)

  # start the search thinking it's ts[1]-ts[2]
  i₋ = 1
  i₊ = 2
  vals = map(idx) do j
    t = tvals[j]

    if continuity === :left
      # we have i₋ = i₊ = 1 if t = ts[1], i₊ = i₋ + 1 = lastindex(ts) if t > ts[end],
      # and otherwise i₋ and i₊ satisfy ts[i₋] < t ≤ ts[i₊]
      i₊ = min(lastindex(ts), _searchsortedfirst(ts,t,i₊,tdir > 0))
      i₋ = i₊ > 1 ? i₊ - 1 : i₊
    else
      # we have i₋ = i₊ - 1 = 1 if t < ts[1], i₊ = i₋ = lastindex(ts) if t = ts[end],
      # and otherwise i₋ and i₊ satisfy ts[i₋] ≤ t < ts[i₊]
      i₋ = max(1, _searchsortedlast(ts,t,i₋,tdir > 0))
      i₊ = i₋ < lastindex(ts) ? i₋ + 1 : i₋
    end

    dt = ts[i₊] - ts[i₋]
    Θ = iszero(dt) ? oneunit(t) / oneunit(dt) : (t-ts[i₋]) / dt

    if typeof(cache) <: (FunctionMapCache) || typeof(cache) <: FunctionMapConstantCache
      return ode_interpolant(Θ,dt,timeseries[i₋],timeseries[i₊],0,cache,idxs,deriv)
    elseif !id.dense
      return linear_interpolant(Θ,dt,timeseries[i₋],timeseries[i₊],idxs,deriv)
    elseif typeof(cache) <: CompositeCache
      DiffEqBase.addsteps!(ks[i₊],ts[i₋],timeseries[i₋],timeseries[i₊],dt,f,p,cache.caches[id.alg_choice[i₊]]) # update the kcurrent
      return ode_interpolant(Θ,dt,timeseries[i₋],timeseries[i₊],ks[i₊],cache.caches[id.alg_choice[i₊]],idxs,deriv)
    else
      DiffEqBase.addsteps!(ks[i₊],ts[i₋],timeseries[i₋],timeseries[i₊],dt,f,p,cache) # update the kcurrent
      return ode_interpolant(Θ,dt,timeseries[i₋],timeseries[i₊],ks[i₊],cache,idxs,deriv)
    end
  end
  invpermute!(vals, idx)
  DiffEqArray(vals, tvals)
end

"""
ode_interpolation(tvals,ts,timeseries,ks)

Get the value at tvals where the solution is known at the
times ts (sorted), with values timeseries and derivatives ks
"""
function ode_interpolation!(vals,tvals,id,idxs,deriv,p,continuity::Symbol=:left)
  @unpack ts,timeseries,ks,f,cache = id
  @inbounds tdir = sign(ts[end]-ts[1])
  idx = sortperm(tvals,rev=tdir<0)

  # start the search thinking it's in ts[1]-ts[2]
  i₋ = 1
  i₊ = 2
  @inbounds for j in idx
    t = tvals[j]

    if continuity === :left
      # we have i₋ = i₊ = 1 if t = ts[1], i₊ = i₋ + 1 = lastindex(ts) if t > ts[end],
      # and otherwise i₋ and i₊ satisfy ts[i₋] < t ≤ ts[i₊]
      i₊ = min(lastindex(ts), _searchsortedfirst(ts,t,i₊,tdir > 0))
      i₋ = i₊ > 1 ? i₊ - 1 : i₊
    else
      # we have i₋ = i₊ - 1 = 1 if t < ts[1], i₊ = i₋ = lastindex(ts) if t = ts[end],
      # and otherwise i₋ and i₊ satisfy ts[i₋] ≤ t < ts[i₊]
      i₋ = max(1, _searchsortedlast(ts,t,i₋,tdir > 0))
      i₊ = i₋ < lastindex(ts) ? i₋ + 1 : i₋
    end

    dt = ts[i₊] - ts[i₋]
    Θ = iszero(dt) ? oneunit(t) / oneunit(dt) : (t-ts[i₋]) / dt

    if typeof(cache) <: (FunctionMapCache) || typeof(cache) <: FunctionMapConstantCache
      if eltype(timeseries) <: AbstractArray
        ode_interpolant!(vals[j],Θ,dt,timeseries[i₋],timeseries[i₊],0,cache,idxs,deriv)
      else
        vals[j] = ode_interpolant(Θ,dt,timeseries[i₋],timeseries[i₊],0,cache,idxs,deriv)
      end
    elseif !id.dense
      if eltype(timeseries) <: AbstractArray
        linear_interpolant!(vals[j],Θ,dt,timeseries[i₋],timeseries[i₊],idxs,deriv)
      else
        vals[j] = linear_interpolant(Θ,dt,timeseries[i₋],timeseries[i₊],idxs,deriv)
      end
    elseif typeof(cache) <: CompositeCache
      DiffEqBase.addsteps!(ks[i₊],ts[i₋],timeseries[i₋],timeseries[i₊],dt,f,p,cache.caches[id.alg_choice[i₊]]) # update the kcurrent
      if eltype(timeseries) <: AbstractArray
        ode_interpolant!(vals[j],Θ,dt,timeseries[i₋],timeseries[i₊],ks[i₊],cache.caches[id.alg_choice[i₊]],idxs,deriv)
      else
        vals[j] = ode_interpolant(Θ,dt,timeseries[i₋],timeseries[i₊],ks[i₊],cache.caches[id.alg_choice[i₊]],idxs,deriv)
      end
    else
      DiffEqBase.addsteps!(ks[i₊],ts[i₋],timeseries[i₋],timeseries[i₊],dt,f,p,cache) # update the kcurrent
      if eltype(vals[j]) <: AbstractArray
        ode_interpolant!(vals[j],Θ,dt,timeseries[i₋],timeseries[i₊],ks[i₊],cache,idxs,deriv)
      else
        vals[j] = ode_interpolant(Θ,dt,timeseries[i₋],timeseries[i₊],ks[i₊],cache,idxs,deriv)
      end
    end
  end

  vals
end

"""
ode_interpolation(tval::Number,ts,timeseries,ks)

Get the value at tval where the solution is known at the
times ts (sorted), with values timeseries and derivatives ks
"""
function ode_interpolation(tval::Number,id,idxs,deriv,p,continuity::Symbol=:left)
  @unpack ts,timeseries,ks,f,cache = id
  @inbounds tdir = sign(ts[end]-ts[1])

  if continuity === :left
    # we have i₋ = i₊ = 1 if tval = ts[1], i₊ = i₋ + 1 = lastindex(ts) if tval > ts[end],
    # and otherwise i₋ and i₊ satisfy ts[i₋] < tval ≤ ts[i₊]
    i₊ = min(lastindex(ts), _searchsortedfirst(ts,tval,2,tdir > 0))
    i₋ = i₊ > 1 ? i₊ - 1 : i₊
  else
    # we have i₋ = i₊ - 1 = 1 if tval < ts[1], i₊ = i₋ = lastindex(ts) if tval = ts[end],
    # and otherwise i₋ and i₊ satisfy ts[i₋] ≤ tval < ts[i₊]
    i₋ = max(1, _searchsortedlast(ts,tval,1,tdir > 0))
    i₊ = i₋ < lastindex(ts) ? i₋ + 1 : i₋
  end

  @inbounds begin
    dt = ts[i₊] - ts[i₋]
    Θ = iszero(dt) ? oneunit(tval) / oneunit(dt) : (tval-ts[i₋]) / dt

    if typeof(cache) <: (FunctionMapCache) || typeof(cache) <: FunctionMapConstantCache
      val = ode_interpolant(Θ,dt,timeseries[i₋],timeseries[i₊],0,cache,idxs,deriv)
    elseif !id.dense
      val = linear_interpolant(Θ,dt,timeseries[i₋],timeseries[i₊],idxs,deriv)
    elseif typeof(cache) <: CompositeCache
      DiffEqBase.addsteps!(ks[i₊],ts[i₋],timeseries[i₋],timeseries[i₊],dt,f,p,cache.caches[id.alg_choice[i₊]]) # update the kcurrent
      val = ode_interpolant(Θ,dt,timeseries[i₋],timeseries[i₊],ks[i₊],cache.caches[id.alg_choice[i₊]],idxs,deriv)
    else
      DiffEqBase.addsteps!(ks[i₊],ts[i₋],timeseries[i₋],timeseries[i₊],dt,f,p,cache) # update the kcurrent
      val = ode_interpolant(Θ,dt,timeseries[i₋],timeseries[i₊],ks[i₊],cache,idxs,deriv)
    end
  end

  val
end

"""
ode_interpolation!(out,tval::Number,ts,timeseries,ks)

Get the value at tval where the solution is known at the
times ts (sorted), with values timeseries and derivatives ks
"""
function ode_interpolation!(out,tval::Number,id,idxs,deriv,p,continuity::Symbol=:left)
  @unpack ts,timeseries,ks,f,cache = id
  @inbounds tdir = sign(ts[end]-ts[1])

  if continuity === :left
    # we have i₋ = i₊ = 1 if tval = ts[1], i₊ = i₋ + 1 = lastindex(ts) if tval > ts[end],
    # and otherwise i₋ and i₊ satisfy ts[i₋] < tval ≤ ts[i₊]
    i₊ = min(lastindex(ts), _searchsortedfirst(ts,tval,2,tdir > 0))
    i₋ = i₊ > 1 ? i₊ - 1 : i₊
  else
    # we have i₋ = i₊ - 1 = 1 if tval < ts[1], i₊ = i₋ = lastindex(ts) if tval = ts[end],
    # and otherwise i₋ and i₊ satisfy ts[i₋] ≤ tval < ts[i₊]
    i₋ = max(1, _searchsortedlast(ts,tval,1,tdir > 0))
    i₊ = i₋ < lastindex(ts) ? i₋ + 1 : i₋
  end

  @inbounds begin
    dt = ts[i₊] - ts[i₋]
    Θ = iszero(dt) ? oneunit(tval) / oneunit(dt) : (tval-ts[i₋]) / dt

    if typeof(cache) <: (FunctionMapCache) || typeof(cache) <: FunctionMapConstantCache
      ode_interpolant!(out,Θ,dt,timeseries[i₋],timeseries[i₊],0,cache,idxs,deriv)
    elseif !id.dense
      linear_interpolant!(out,Θ,dt,timeseries[i₋],timeseries[i₊],idxs,deriv)
    elseif typeof(cache) <: CompositeCache
      DiffEqBase.addsteps!(ks[i₊],ts[i₋],timeseries[i₋],timeseries[i₊],dt,f,p,cache.caches[id.alg_choice[i₊]]) # update the kcurrent
      ode_interpolant!(out,Θ,dt,timeseries[i₋],timeseries[i₊],ks[i₊],cache.caches[id.alg_choice[i₊]],idxs,deriv)
    else
      DiffEqBase.addsteps!(ks[i₊],ts[i₋],timeseries[i₋],timeseries[i₊],dt,f,p,cache) # update the kcurrent
      ode_interpolant!(out,Θ,dt,timeseries[i₋],timeseries[i₊],ks[i₊],cache,idxs,deriv)
    end
  end

  out
end

"""
By default, Hermite interpolant so update the derivative at the two ends
"""
function DiffEqBase.addsteps!(k,t,uprev,u,dt,f,p,cache,always_calc_begin = false,allow_calc_end = true,force_calc_end = false)
  if length(k)<2 || always_calc_begin
    if typeof(cache) <: OrdinaryDiffEqMutableCache
      rtmp = similar(u, eltype(eltype(k)))
      f(rtmp,uprev,p,t)
      copyat_or_push!(k,1,rtmp)
      f(rtmp,u,p,t+dt)
      copyat_or_push!(k,2,rtmp)
    else
      copyat_or_push!(k,1,f(uprev,p,t))
      copyat_or_push!(k,2,f(u,p,t+dt))
    end
  end
  nothing
end

"""
ode_interpolant and ode_interpolant! dispatch
"""
function ode_interpolant(Θ,dt,y₀,y₁,k,cache,idxs,T::Type{Val{TI}}) where TI
  _ode_interpolant(Θ,dt,y₀,y₁,k,cache,idxs,T)
end

function ode_interpolant(Θ,dt,y₀,y₁,k,cache::OrdinaryDiffEqMutableCache,idxs,T::Type{Val{TI}}) where TI
  if typeof(idxs) <: Number || typeof(y₀) <: Union{Number,SArray}
    # typeof(y₀) can be these if saveidxs gives a single value
    _ode_interpolant(Θ,dt,y₀,y₁,k,cache,idxs,T)
  elseif typeof(idxs) <: Nothing
    out = oneunit(Θ) .* y₁
    _ode_interpolant!(out,Θ,dt,y₀,y₁,k,cache,idxs,T)
  else
    out = oneunit(Θ) .* y₁[idxs]
    _ode_interpolant!(out,Θ,dt,y₀,y₁,k,cache,idxs,T)
  end
end

function ode_interpolant!(out,Θ,dt,y₀,y₁,k,cache,idxs,T::Type{Val{TI}}) where TI
  _ode_interpolant!(out,Θ,dt,y₀,y₁,k,cache,idxs,T)
end

##################### Hermite Interpolants

# If no dispatch found, assume Hermite
function _ode_interpolant(Θ,dt,y₀,y₁,k,cache,idxs,T::Type{Val{TI}}) where TI
  hermite_interpolant(Θ,dt,y₀,y₁,k,Val{typeof(cache) <: OrdinaryDiffEqMutableCache},idxs,T)
end

function _ode_interpolant!(out,Θ,dt,y₀,y₁,k,cache,idxs,T::Type{Val{TI}}) where TI
  hermite_interpolant!(out,Θ,dt,y₀,y₁,k,idxs,T)
end

"""
Hairer Norsett Wanner Solving Ordinary Differential Euations I - Nonstiff Problems Page 190

Herimte Interpolation, chosen if no other dispatch for ode_interpolant
"""
@muladd function hermite_interpolant(Θ,dt,y₀,y₁,k,::Type{Val{false}},idxs::Nothing,T::Type{Val{0}}) # Default interpolant is Hermite
  #@.. (1-Θ)*y₀+Θ*y₁+Θ*(Θ-1)*((1-2Θ)*(y₁-y₀)+(Θ-1)*dt*k[1] + Θ*dt*k[2])
  @inbounds (1-Θ)*y₀+Θ*y₁+Θ*(Θ-1)*((1-2Θ)*(y₁-y₀)+(Θ-1)*dt*k[1] + Θ*dt*k[2])
end

@muladd function hermite_interpolant(Θ,dt,y₀,y₁,k,::Type{Val{true}},idxs::Nothing,T::Type{Val{0}}) # Default interpolant is Hermite
  #@.. (1-Θ)*y₀+Θ*y₁+Θ*(Θ-1)*((1-2Θ)*(y₁-y₀)+(Θ-1)*dt*k[1] + Θ*dt*k[2])
  @inbounds @.. (1-Θ)*y₀+Θ*y₁+Θ*(Θ-1)*((1-2Θ)*(y₁-y₀)+(Θ-1)*dt*k[1] + Θ*dt*k[2])
end

@muladd function hermite_interpolant(Θ,dt,y₀,y₁,k,cache,idxs,T::Type{Val{0}}) # Default interpolant is Hermite
  # return @.. (1-Θ)*y₀[idxs]+Θ*y₁[idxs]+Θ*(Θ-1)*((1-2Θ)*(y₁[idxs]-y₀[idxs])+(Θ-1)*dt*k[1][idxs] + Θ*dt*k[2][idxs])
  return (1-Θ)*y₀[idxs]+Θ*y₁[idxs]+Θ*(Θ-1)*((1-2Θ)*(y₁[idxs]-y₀[idxs])+(Θ-1)*dt*k[1][idxs] + Θ*dt*k[2][idxs])
end

@muladd function hermite_interpolant!(out,Θ,dt,y₀,y₁,k,idxs::Nothing,T::Type{Val{0}}) # Default interpolant is Hermite
  @inbounds @.. out = (1-Θ)*y₀+Θ*y₁+Θ*(Θ-1)*((1-2Θ)*(y₁-y₀)+(Θ-1)*dt*k[1] + Θ*dt*k[2])
  #@inbounds for i in eachindex(out)
  #  out[i] = (1-Θ)*y₀[i]+Θ*y₁[i]+Θ*(Θ-1)*((1-2Θ)*(y₁[i]-y₀[i])+(Θ-1)*dt*k[1][i] + Θ*dt*k[2][i])
  #end
  #out
end

@muladd function hermite_interpolant!(out,Θ,dt,y₀,y₁,k,idxs,T::Type{Val{0}}) # Default interpolant is Hermite
  @views @.. out = (1-Θ)*y₀[idxs]+Θ*y₁[idxs]+Θ*(Θ-1)*((1-2Θ)*(y₁[idxs]-y₀[idxs])+(Θ-1)*dt*k[1][idxs] + Θ*dt*k[2][idxs])
  #@inbounds for (j,i) in enumerate(idxs)
  #  out[j] = (1-Θ)*y₀[i]+Θ*y₁[i]+Θ*(Θ-1)*((1-2Θ)*(y₁[i]-y₀[i])+(Θ-1)*dt*k[1][i] + Θ*dt*k[2][i])
  #end
  #out
end

"""
Herimte Interpolation, chosen if no other dispatch for ode_interpolant
"""
@muladd function hermite_interpolant(Θ,dt,y₀,y₁,k,::Type{Val{false}},idxs::Nothing,T::Type{Val{1}}) # Default interpolant is Hermite
  #@.. k[1] + Θ*(-4*dt*k[1] - 2*dt*k[2] - 6*y₀ + Θ*(3*dt*k[1] + 3*dt*k[2] + 6*y₀ - 6*y₁) + 6*y₁)/dt
  @inbounds k[1] + Θ*(-4*dt*k[1] - 2*dt*k[2] - 6*y₀ + Θ*(3*dt*k[1] + 3*dt*k[2] + 6*y₀ - 6*y₁) + 6*y₁)/dt
end

@muladd function hermite_interpolant(Θ,dt,y₀,y₁,k,::Type{Val{true}},idxs::Nothing,T::Type{Val{1}}) # Default interpolant is Hermite
  @inbounds @.. k[1] + Θ*(-4*dt*k[1] - 2*dt*k[2] - 6*y₀ + Θ*(3*dt*k[1] + 3*dt*k[2] + 6*y₀ - 6*y₁) + 6*y₁)/dt
end

@muladd function hermite_interpolant(Θ,dt,y₀,y₁,k,cache,idxs,T::Type{Val{1}}) # Default interpolant is Hermite
  # return @.. k[1][idxs] + Θ*(-4*dt*k[1][idxs] - 2*dt*k[2][idxs] - 6*y₀[idxs] + Θ*(3*dt*k[1][idxs] + 3*dt*k[2][idxs] + 6*y₀[idxs] - 6*y₁[idxs]) + 6*y₁[idxs])/dt
  return k[1][idxs] + Θ*(-4*dt*k[1][idxs] - 2*dt*k[2][idxs] - 6*y₀[idxs] + Θ*(3*dt*k[1][idxs] + 3*dt*k[2][idxs] + 6*y₀[idxs] - 6*y₁[idxs]) + 6*y₁[idxs])/dt
end

@muladd function hermite_interpolant!(out,Θ,dt,y₀,y₁,k,idxs::Nothing,T::Type{Val{1}}) # Default interpolant is Hermite
  @inbounds @.. out = k[1] + Θ*(-4*dt*k[1] - 2*dt*k[2] - 6*y₀ + Θ*(3*dt*k[1] + 3*dt*k[2] + 6*y₀ - 6*y₁) + 6*y₁)/dt
  #@inbounds for i in eachindex(out)
  #  out[i] = k[1][i] + Θ*(-4*dt*k[1][i] - 2*dt*k[2][i] - 6*y₀[i] + Θ*(3*dt*k[1][i] + 3*dt*k[2][i] + 6*y₀[i] - 6*y₁[i]) + 6*y₁[i])/dt
  #end
  #out
end

@muladd function hermite_interpolant!(out,Θ,dt,y₀,y₁,k,idxs,T::Type{Val{1}}) # Default interpolant is Hermite
  @views @.. out = k[1][idxs] + Θ*(-4*dt*k[1][idxs] - 2*dt*k[2][idxs] - 6*y₀[idxs] + Θ*(3*dt*k[1][idxs] + 3*dt*k[2][idxs] + 6*y₀[idxs] - 6*y₁[idxs]) + 6*y₁[idxs])/dt
  #@inbounds for (j,i) in enumerate(idxs)
  #  out[j] = k[1][i] + Θ*(-4*dt*k[1][i] - 2*dt*k[2][i] - 6*y₀[i] + Θ*(3*dt*k[1][i] + 3*dt*k[2][i] + 6*y₀[i] - 6*y₁[i]) + 6*y₁[i])/dt
  #end
  #out
end

"""
Herimte Interpolation, chosen if no other dispatch for ode_interpolant
"""
@muladd function hermite_interpolant(Θ,dt,y₀,y₁,k,::Type{Val{false}},idxs::Nothing,T::Type{Val{2}}) # Default interpolant is Hermite
  #@.. (-4*dt*k[1] - 2*dt*k[2] - 6*y₀ + Θ*(6*dt*k[1] + 6*dt*k[2] + 12*y₀ - 12*y₁) + 6*y₁)/(dt*dt)
  @inbounds (-4*dt*k[1] - 2*dt*k[2] - 6*y₀ + Θ*(6*dt*k[1] + 6*dt*k[2] + 12*y₀ - 12*y₁) + 6*y₁)/(dt*dt)
end

@muladd function hermite_interpolant(Θ,dt,y₀,y₁,k,::Type{Val{true}},idxs::Nothing,T::Type{Val{2}}) # Default interpolant is Hermite
  #@.. (-4*dt*k[1] - 2*dt*k[2] - 6*y₀ + Θ*(6*dt*k[1] + 6*dt*k[2] + 12*y₀ - 12*y₁) + 6*y₁)/(dt*dt)
  @inbounds @.. (-4*dt*k[1] - 2*dt*k[2] - 6*y₀ + Θ*(6*dt*k[1] + 6*dt*k[2] + 12*y₀ - 12*y₁) + 6*y₁)/(dt*dt)
end

@muladd function hermite_interpolant(Θ,dt,y₀,y₁,k,cache,idxs,T::Type{Val{2}}) # Default interpolant is Hermite
  #out = similar(y₀,axes(idxs))
  #@views @.. out = (-4*dt*k[1][idxs] - 2*dt*k[2][idxs] - 6*y₀[idxs] + Θ*(6*dt*k[1][idxs] + 6*dt*k[2][idxs] + 12*y₀[idxs] - 12*y₁[idxs]) + 6*y₁[idxs])/(dt*dt)
  @views out = (-4*dt*k[1][idxs] - 2*dt*k[2][idxs] - 6*y₀[idxs] + Θ*(6*dt*k[1][idxs] + 6*dt*k[2][idxs] + 12*y₀[idxs] - 12*y₁[idxs]) + 6*y₁[idxs])/(dt*dt)
  out
end

@muladd function hermite_interpolant!(out,Θ,dt,y₀,y₁,k,idxs::Nothing,T::Type{Val{2}}) # Default interpolant is Hermite
  @inbounds @.. out = (-4*dt*k[1] - 2*dt*k[2] - 6*y₀ + Θ*(6*dt*k[1] + 6*dt*k[2] + 12*y₀ - 12*y₁) + 6*y₁)/(dt*dt)
  #@inbounds for i in eachindex(out)
  #  out[i] = (-4*dt*k[1][i] - 2*dt*k[2][i] - 6*y₀[i] + Θ*(6*dt*k[1][i] + 6*dt*k[2][i] + 12*y₀[i] - 12*y₁[i]) + 6*y₁[i])/(dt*dt)
  #end
  #out
end

@muladd function hermite_interpolant!(out,Θ,dt,y₀,y₁,k,idxs,T::Type{Val{2}}) # Default interpolant is Hermite
  @views @.. out = (-4*dt*k[1][idxs] - 2*dt*k[2][idxs] - 6*y₀[idxs] + Θ*(6*dt*k[1][idxs] + 6*dt*k[2][idxs] + 12*y₀[idxs] - 12*y₁[idxs]) + 6*y₁[idxs])/(dt*dt)
  #@inbounds for (j,i) in enumerate(idxs)
  #  out[j] = (-4*dt*k[1][i] - 2*dt*k[2][i] - 6*y₀[i] + Θ*(6*dt*k[1][i] + 6*dt*k[2][i] + 12*y₀[i] - 12*y₁[i]) + 6*y₁[i])/(dt*dt)
  #end
  #out
end

"""
Herimte Interpolation, chosen if no other dispatch for ode_interpolant
"""
@muladd function hermite_interpolant(Θ,dt,y₀,y₁,k,::Type{Val{false}},idxs::Nothing,T::Type{Val{3}}) # Default interpolant is Hermite
  #@.. (6*dt*k[1] + 6*dt*k[2] + 12*y₀ - 12*y₁)/(dt*dt*dt)
  @inbounds (6*dt*k[1] + 6*dt*k[2] + 12*y₀ - 12*y₁)/(dt*dt*dt)
end

@muladd function hermite_interpolant(Θ,dt,y₀,y₁,k,::Type{Val{true}},idxs::Nothing,T::Type{Val{3}}) # Default interpolant is Hermite
  #@.. (6*dt*k[1] + 6*dt*k[2] + 12*y₀ - 12*y₁)/(dt*dt*dt)
  @inbounds @.. (6*dt*k[1] + 6*dt*k[2] + 12*y₀ - 12*y₁)/(dt*dt*dt)
end

@muladd function hermite_interpolant(Θ,dt,y₀,y₁,k,cache,idxs,T::Type{Val{3}}) # Default interpolant is Hermite
  #out = similar(y₀,axes(idxs))
  #@views @.. out = (6*dt*k[1][idxs] + 6*dt*k[2][idxs] + 12*y₀[idxs] - 12*y₁[idxs])/(dt*dt*dt)
  @views out = (6*dt*k[1][idxs] + 6*dt*k[2][idxs] + 12*y₀[idxs] - 12*y₁[idxs])/(dt*dt*dt)
  out
end

@muladd function hermite_interpolant!(out,Θ,dt,y₀,y₁,k,idxs::Nothing,T::Type{Val{3}}) # Default interpolant is Hermite
  @inbounds @.. out = (6*dt*k[1] + 6*dt*k[2] + 12*y₀ - 12*y₁)/(dt*dt*dt)
  #for i in eachindex(out)
  #  out[i] = (6*dt*k[1][i] + 6*dt*k[2][i] + 12*y₀[i] - 12*y₁[i])/(dt*dt*dt)
  #end
  #out
end

@muladd function hermite_interpolant!(out,Θ,dt,y₀,y₁,k,idxs,T::Type{Val{3}}) # Default interpolant is Hermite
  @views @.. out = (6*dt*k[1][idxs] + 6*dt*k[2][idxs] + 12*y₀[idxs] - 12*y₁[idxs])/(dt*dt*dt)
  #for (j,i) in enumerate(idxs)
  #  out[j] = (6*dt*k[1][i] + 6*dt*k[2][i] + 12*y₀[i] - 12*y₁[i])/(dt*dt*dt)
  #end
  #out
end

######################## Linear Interpolants

@muladd @inline function linear_interpolant(Θ,dt,y₀,y₁,idxs::Nothing,T::Type{Val{0}})
  Θm1 = (1-Θ)
  @.. Θm1*y₀ + Θ*y₁
end

@muladd @inline function linear_interpolant(Θ,dt,y₀,y₁,idxs,T::Type{Val{0}})
  Θm1 = (1-Θ)
  @.. Θm1*y₀[idxs] + Θ*y₁[idxs]
end

@muladd @inline function linear_interpolant!(out,Θ,dt,y₀,y₁,idxs::Nothing,T::Type{Val{0}})
  Θm1 = (1-Θ)
  @.. out = Θm1*y₀ + Θ*y₁
  out
end

@muladd @inline function linear_interpolant!(out,Θ,dt,y₀,y₁,idxs,T::Type{Val{0}})
  Θm1 = (1-Θ)
  @views @.. out = Θm1*y₀[idxs] + Θ*y₁[idxs]
  out
end

"""
Linear Interpolation
"""
@inline function linear_interpolant(Θ,dt,y₀,y₁,idxs::Nothing,T::Type{Val{1}})
  (y₁ - y₀)/dt
end

@inline function linear_interpolant(Θ,dt,y₀,y₁,idxs,T::Type{Val{1}})
  @.. (y₁[idxs] - y₀[idxs])/dt
end

@inline function linear_interpolant!(out,Θ,dt,y₀,y₁,idxs::Nothing,T::Type{Val{1}})
  @.. out = (y₁ - y₀)/dt
  out
end

@inline function linear_interpolant!(out,Θ,dt,y₀,y₁,idxs,T::Type{Val{1}})
  @views @.. out = (y₁[idxs] - y₀[idxs])/dt
  out
end


"""测试不同的粒度，bubble的区别"""

using Plots
using MARY_fRG.Basics
using MARY_fRG.Drawers
using MARY_fRG.Fermi
using MARY_fRG.Fermi.Patch
using MARY_fRG.Fermi.Kagome
using MARY_fRG.FlowEquation

"""筛选tris"""
function fliter(Γ4::Gamma4{T, P}, lval) where {T, P}
    lamb = Γ4.λ_0 * exp(-lval)
    resltris::Vector{P} = []
    reslpats::Vector{Int64} = []
    cert = 1#length(Γ4.ltris) > 2000000 ? 5 : 15
    mintri = missing
    minpat = missing
    mineng = 10.
    for (tri, pat) in zip(Γ4.ltris, Γ4.lpats)
        if pat != 1
            continue
        end
        engs = [Γ4.model.dispersion[bidx](
            tri.center.x, tri.center.y) / lamb for bidx in 1:1:Γ4.model.bandnum
        ]
        #如果还有可能有贡献
        if minimum(abs.(engs)) < cert #&& minimum(abs.(engs2)) > cert # && pat == 1# && minimum(abs.(engs)) > 0.1
            push!(resltris, tri)
            push!(reslpats, pat)
        end
        #if minimum(abs.(engs)) < mineng
        #    mintri = tri
        #    minpat = pat
        #    mineng = minimum(abs.(engs))
        #end
    end
    #
    #resltris = [mintri]#resltris[1:1]
    #reslpats = [minpat]#reslpats[1:1]
    #
    ltris_pat = Vector{typeof(resltris)}(undef, Γ4.patchnum)
    for idx in 1:1:Γ4.patchnum
        ltris_pat[idx] = []
    end
    #
    for (tri, pat) in zip(resltris, reslpats)
        push!(ltris_pat[pat], tri)
    end
    #
    return Gamma4(
        Γ4.model,
        Γ4.λ_0,
        Γ4.V,
        Γ4.k4tab,
        Γ4.patchnum,
        Γ4.patches,
        resltris, reslpats, nothing, ltris_pat
    )
end

function refine(Γ4::Gamma4{T, P}) where {T, P}
    resltris::Vector{P} = []
    reslpats::Vector{Int64} = []
    resltris = refine_list_triangle(Γ4.ltris, 5)
    reslpats = group_ltris_into_patches_mt(resltris, Γ4.model.brillouin, Γ4.patchnum)
    #
    #resltris = [mintri]#resltris[1:1]
    #reslpats = [minpat]#reslpats[1:1]
    #
    ltris_pat = Vector{typeof(resltris)}(undef, Γ4.patchnum)
    for idx in 1:1:Γ4.patchnum
        ltris_pat[idx] = []
    end
    #
    for (tri, pat) in zip(resltris, reslpats)
        push!(ltris_pat[pat], tri)
    end
    #
    return Gamma4(
        Γ4.model,
        Γ4.λ_0,
        Γ4.V,
        Γ4.k4tab,
        Γ4.patchnum,
        Γ4.patches,
        resltris, reslpats, nothing, ltris_pat
    )
end

function run_rftf()
    model = common_square_lattice(0.20)
    disp = model.dispersion[1]
    lval = 15.
    #500
    Γ4_5 = TFGamma4(
        model, 8.0, 16, 500
    )
    Γ4_5.V .+= 1.0
    Γ4_5_1 = fliter(Γ4_5, 11.0)
    Γ4_5_2 = fliter(Γ4_5, 12.0)
    println(length(Γ4_5_1.ltris))
    println(length(Γ4_5_2.ltris))
    #fliter2
    bubbval_fs = Array{Float64, 2}(undef, Γ4_5_2.patchnum, Γ4_5_2.patchnum)
    lamb = Γ4_5_2.λ_0 * exp(-lval)
    println("lamb ", lamb)
    Threads.@threads for idxs in CartesianIndices(bubbval_fs)
        i2, i3 = Tuple(idxs)
        k2, k3 = Γ4_5_2.patches[1][i2], Γ4_5_2.patches[1][i3]
        q_fs = Γ4_5_2.model.kadd(k3, -k2)
        bubbres = pi_αβ_plus_tf(
            Γ4_5_2.ltris, area(Γ4_5_2.model.brillouin),
            lamb, q_fs, Γ4_5_2.model.dispersion[1],
            Γ4_5_2.model.dispersion[1], Γ4_5_2.model.kadd
        )
        bubbval_fs[i2, i3] = bubbres
    end
    plt = heatmap(1e5*bubbval_fs)
    @savepng plt "bubb_5_2_2"
    #fliter1
    bubbval_fs = Array{Float64, 2}(undef, Γ4_5_1.patchnum, Γ4_5_1.patchnum)
    lamb = Γ4_5_1.λ_0 * exp(-lval)
    Threads.@threads for idxs in CartesianIndices(bubbval_fs)
        i2, i3 = Tuple(idxs)
        k2, k3 = Γ4_5_1.patches[1][i2], Γ4_5_1.patches[1][i3]
        q_fs = Γ4_5_1.model.kadd(k3, -k2)
        bubbres = pi_αβ_plus_tf(
            Γ4_5_1.ltris, area(Γ4_5_1.model.brillouin),
            lamb, q_fs, Γ4_5_1.model.dispersion[1],
            Γ4_5_1.model.dispersion[1], Γ4_5_1.model.kadd
        )
        bubbval_fs[i2, i3] = bubbres
    end
    plt = heatmap(1e5*bubbval_fs)
    @savepng plt "bubb_5_2_1"
    #fliter3
    ltris3::Vector{Basics.AbstractTriangle{:RT}} = []
    ltris4::Vector{Basics.AbstractTriangle{:RT}} = []
    for tri in Γ4_5_1.ltris
        if !any(map((x)->x==tri, Γ4_5_2.ltris))
            push!(ltris4, tri)
        end
    end
    #push!(ltris3, Γ4_5_2.ltris[1])
    #println(Γ4_5_2.ltris[1])
    #println(disp(Γ4_5_2.ltris[1].center.x, Γ4_5_2.ltris[1].center.y))
    #push!(ltris3, Γ4_5_2.ltris[2])
    push!(ltris3, Γ4_5_2.ltris[3])
    #push!(ltris3, ltris4[1])
    push!(ltris3, ltris4[2])
    println(disp(ltris4[2].center.x, ltris4[2].center.y))
    #push!(ltris3, ltris4[3])
    println(length(ltris3))
    bubbval_fs = Array{Float64, 2}(undef, Γ4_5_1.patchnum, Γ4_5_1.patchnum)
    lamb = Γ4_5_1.λ_0 * exp(-lval)
    Threads.@threads for idxs in CartesianIndices(bubbval_fs)
        i2, i3 = Tuple(idxs)
        k2, k3 = Γ4_5_1.patches[1][i2], Γ4_5_1.patches[1][i3]
        q_fs = Γ4_5_1.model.kadd(k3, -k2)
        bubbres = pi_αβ_plus_tf(
            ltris3, area(Γ4_5_1.model.brillouin),
            lamb, q_fs, Γ4_5_1.model.dispersion[1],
            Γ4_5_1.model.dispersion[1], Γ4_5_1.model.kadd
        )
        bubbval_fs[i2, i3] = bubbres
    end
    plt = heatmap(1e5*bubbval_fs)
    @savepng plt "bubb_5_2_3"
    return
    #plt = draw_points([tri.center for tri in Γ4_5.ltris])
    #abses = [abs(disp(tri.center.x, tri.center.y)) for tri in Γ4_5.ltris]
    #println(minimum(abses))
    #abscount = 0
    #for ae in abses
    #    if ae < 0.01
    #        abscount += 1
    #    end
    #end
    #println(abscount)
    #@savepng plt "diffe1"
    println(Γ4_5.ltris[1])
    bubb_pp_5, bubb_fs_5, bubb_ex_5 = all_bubble_tf_mt(Γ4_5, lval)
    println("500时的贡献 ", sum(abs.(bubb_fs_5.V)))
    println("500时最大", findmax(abs.(bubb_fs_5.V)))
    k1, k2 = Γ4_5.patches[1][12], Γ4_5.patches[1][4]
    println(k1)
    println(k2)
    kprim = Γ4_5.model.kadd(k1, k2)
    println(kprim)
    kprim = Γ4_5.model.kadd(kprim, -Γ4_5.ltris[1].center)
    println(kprim)
    println("kprim disp ", Γ4_5.model.dispersion[1](kprim.x, kprim.y))

    return
    ###
    #100
    Γ4_1 = TFGamma4(
        model, 8.0, 16, 800
    )
    Γ4_1.V .+= 1.0
    #Γ4_1 = fliter_away_surface(Γ4_1, 10.0)
    #plt = draw_points([tri.center for tri in Γ4_1.ltris])
    #@savepng plt "diffe3"
    #Γ4_1 = TFGamma4_addition_ltris_mt(Γ4_1, 10.0)
    Γ4_1 = fliter(Γ4_1, 12.0)
    #println(Γ4_1.ltris)
    #plt = draw_polygon(Γ4_1.ltris)
    #plot!(plt, aspect_ratio=:equal)
    #@savepng plt "diffe4"
    Γ4_1 = refine(Γ4_1)
    #println(Γ4_1.ltris)
    plt = draw_polygon(Γ4_1.ltris)
    #plot!(plt, aspect_ratio=:equal)
    #@savepng plt "diffe5"
    Γ4_1 = fliter(Γ4_1, 12.0)
    abses = [abs(disp(tri.center.x, tri.center.y)) for tri in Γ4_1.ltris]
    println(minimum(abses))
    abscount = 0
    for ae in abses
        if ae < 0.01
            abscount += 1
        end
    end
    println(abscount)
    #println(minimum([abs(disp(tri.center.x, tri.center.y)) for tri in Γ4_1.ltris]))
    plt = draw_points([tri.center for tri in Γ4_1.ltris])
    @savepng plt "diffe2"
    println(Γ4_1.ltris[1])
    k1, k2 = Γ4_1.patches[1][11], Γ4_1.patches[1][3]
    println(k1)
    println(k2)
    kprim = Γ4_1.model.kadd(k1, k2)
    println(kprim)
    kprim = Γ4_1.model.kadd(kprim, -Γ4_1.ltris[1].center)
    println(kprim)
    println("kprim disp ", Γ4_1.model.dispersion[1](kprim.x, kprim.y))
    bubb_pp_1, bubb_fs_1, bubb_ex_1 = all_bubble_tf_mt(Γ4_1, lval)
    println("200时的贡献 ", sum(abs.(bubb_pp_1.V)))
    println("200时最大", findmax(abs.(bubb_pp_1.V)))
    #500和100的差
    diffe = bubb_pp_5.V - bubb_pp_1.V
    diffe = abs.(diffe)
    println(sum(diffe))
    println(sum(abs.(bubb_pp_1.V)))
    #refine
    #Γ4_r = TFGamma4_addition_ltris_mt(Γ4_1, lval)
    #bubb_pp_r, bubb_fs_r, bubb_ex_r = all_bubble_tf_mt(Γ4_r, lval)
    ##500和100的差
    #diffe = bubb_pp_5.V - bubb_pp_r.V
    #diffe = abs.(diffe)
    #println(sum(diffe))
end

#run_rftf()

function run_rf_ult()
    #model = common_square_lattice(0.20)
    model = upperband_kagome_lattice(-0.2)
    disp = model.dispersion[1]
    lval = 17.
    #500
    Γ4_5 = TFGamma4(
        model, 8.0, 24, 500
    )
    Γ4_5.V .+= 1.0
    #Γ4_5 = fliter(Γ4_5, 12.0)
    #plt = draw_points([tri.center for tri in Γ4_5.ltris])
    #abses = [abs(disp(tri.center.x, tri.center.y)) for tri in Γ4_5.ltris]
    #println(minimum(abses))
    #abscount = 0
    #for ae in abses
    #    if ae < 0.01
    #        abscount += 1
    #    end
    #end
    #println(abscount)
    #@savepng plt "diffe1"
    println(Γ4_5.ltris[1])
    bubb_pp_5, bubb_fs_5, bubb_ex_5 = all_bubble_tf_mt(Γ4_5, lval)
    println("500时的贡献 ", sum(abs.(bubb_pp_5.V)))
    println("500时最大", findmax(abs.(bubb_pp_5.V)))
    ###
    #100
    Γ4_1 = TFGamma4(
        model, 8.0, 24, 200
    )
    Γ4_1.V .+= 1.0
    bubb_pp_1, bubb_fs_1, bubb_ex_1 = all_bubble_tf_mt(Γ4_1, lval)
    println("200时的贡献 ", sum(abs.(bubb_pp_1.V)))
    println("200时最大", findmax(abs.(bubb_pp_1.V)))
    ###
    #ult
    Γ4_1 = fliter_away_surface(Γ4_1, 6.0)
    Γ4_1 = refine_to_surface(Γ4_1)
    while length(Γ4_1.ltris) < 1000000
        Γ4_1 = refine_to_surface(Γ4_1)
    end
    plt = draw_points([tri.center for tri in Γ4_1.ltris])
    @savepng plt "diffe6"
    println("tri大小", length(Γ4_1.ltris))
    maxi, mini = engpeak_to_surface(Γ4_1)
    println(maxi, " ", mini)
    bubb_pp_u, bubb_fs_u, bubb_ex_u = all_bubble_tf_ult_mt(Γ4_1, maxi)
    println("ult时的贡献 ", sum(abs.(bubb_pp_u.V)))
    println("ult时最大", findmax(abs.(bubb_pp_u.V)))
    #500和100的差
    diffe = bubb_pp_5.V - bubb_pp_1.V
    diffe = abs.(diffe)
    println(sum(diffe))
    println(sum(abs.(bubb_pp_1.V)))
    #500个ult的差
    diffe = bubb_pp_5.V - bubb_pp_u.V
    diffe = abs.(diffe)
    println(sum(diffe))
    println(sum(abs.(bubb_pp_u.V)))
    #ult和500的画图
    plt = heatmap(1e5*bubb_pp_u.V[1, 1, 1, 1, 1, :, :])
    png(plt, "bubb_u1")
    plt = heatmap(1e5*bubb_pp_u.V[1, 1, 1, 1, 2, :, :])
    png(plt, "bubb_u2")
    plt = heatmap(1e5*bubb_pp_u.V[1, 1, 1, 1, 3, :, :])
    png(plt, "bubb_u3")
    plt = heatmap(1e5*bubb_pp_u.V[1, 1, 1, 1, 4, :, :])
    png(plt, "bubb_u4")
    plt = heatmap(1e5*bubb_pp_5.V[1, 1, 1, 1, 1, :, :])
    png(plt, "bubb_51")
    plt = heatmap(1e5*bubb_pp_5.V[1, 1, 1, 1, 2, :, :])
    png(plt, "bubb_52")
    plt = heatmap(1e5*bubb_pp_5.V[1, 1, 1, 1, 3, :, :])
    png(plt, "bubb_53")
    plt = heatmap(1e5*bubb_pp_5.V[1, 1, 1, 1, 4, :, :])
    png(plt, "bubb_54")
    println(Γ4_1.patches[1][1], Γ4_1.patches[1][4], Γ4_1.patches[1][9])
end

run_rf_ult()


using StructArrays
using ChainRulesCore: NoTangent

struct ∇getindex{T,S}
    xs::T
    i::S
end

function (g::∇getindex)(Δ)
    Δ′ = zero(g.xs)
    Δ′[g.i...] = Δ
    (ChainRulesCore.NoTangent(), Δ′, map(_ -> nothing, g.i)...)
end

function ChainRulesCore.rrule(g::∇getindex, Δ)
    g(Δ), Δ′′->(nothing, Δ′′[1][g.i...])
end

function ChainRulesCore.rrule(::typeof(getindex), xs::Array, i...)
    xs[i...], ∇getindex(xs, i)
end

function reversediff(f, xs...)
    y, f☆ = ∂⃖(f, xs...)
    return tuple(y, tail(f☆(dx(y)))...)
end

function reversediff_array(f, xs::Vector...)
    fieldarrays(StructArray(reversediff(f, x...) for x in zip(xs...)))
end

function reversediff_array(f, xs::Vector)
    fieldarrays(StructArray(reversediff(f, x) for x in xs))
end

function assert_gf(f)
    @assert sizeof(sin) == 0
end

function ChainRulesCore.rrule(::typeof(assert_gf), f)
    assert_gf(f), Δ->begin
        (NoTangent(), NoTangent())
    end
end

#=
function ChainRulesCore.rrule(::typeof(map), f, xs::Vector...)
    assert_gf(f)
    primal, dual = reversediff_array(f, xs...)
    primal, Δ->begin
        (NoTangent(), NoTangent(), ntuple(i->map(*, getfield(dual, i), Δ), length(dual))...)
    end
end
=#

# Disable thunk versions of ChainRules, which interfere with higher order AD

function rrule_times(::typeof(*), A::AbstractVecOrMat, B::AbstractVecOrMat)
    function times_pullback(Ȳ)
        return (NoTangent(), Ȳ * Base.adjoint(B), Base.adjoint(A) * Ȳ)
    end
    return A * B, times_pullback
end

function rrule_times(::typeof(*), A::AbstractVector{<:ChainRules.CommutativeMulNumber}, B::AbstractMatrix{<:ChainRules.CommutativeMulNumber})
    function times_pullback(Ȳ)
        return (NoTangent(), Ȳ * Base.adjoint(B), Base.adjoint(A) * Ȳ)
    end
    return A * B, times_pullback
end

rrule_times(::typeof(*), args...) = rrule(*, args...)

function (::∂⃖{N})(f::typeof(*), args...) where {N}
    if N == 1
        z = rrule_times(f, args...)
        if z === nothing
            return ∂⃖recurse{1}()(f, args...)
        end
        return z
    else
        ∂⃖p = ∂⃖{minus1(N)}()
        @destruct z, z̄ = ∂⃖p(rrule_times, f, args...)
        if z === nothing
            return ∂⃖recurse{N}()(f, args...)
        else
            return ∂⃖rrule{N}()(z, z̄)
        end
    end
end

function ChainRulesCore.frule((_, ∂A, ∂B), ::typeof(*), A::AbstractMatrix{<:Real}, B::AbstractVector{<:Real})
    return (A * B, ∂A * B + A * ∂B)
end

#=
function ChainRulesCore.rrule(::typeof(map), f, xs::Vector)
    assert_gf(f)
    arrs = reversediff_array(f, xs)
    primal = getfield(arrs, 1)
    primal, let dual = getfield(arrs, 2)
        Δ->(NoTangent(), NoTangent(), map(*, dual, unthunk(Δ)))
    end
end
=#

#=
function ChainRulesCore.rrule(::typeof(map), f, xs::Vector, ys::Vector)
    assert_gf(f)
    arrs = reversediff_array(f, xs, ys)
    primal = getfield(arrs, 1)
    primal, let dual = tail(arrs)
        Δ->(NoTangent(), NoTangent(), map(*, getfield(dual, 1), Δ), map(*, getfield(dual, 2), Δ))
    end
end
=#

xsum(x::Vector) = sum(x)
function ChainRulesCore.rrule(::typeof(xsum), x::Vector)
    xsum(x), let xdims=size(x)
        Δ->(NoTangent(), xfill(Δ, xdims...))
    end
end

xfill(x, dims...) = fill(x, dims...)
function ChainRulesCore.rrule(::typeof(xfill), x, dim)
    xfill(x, dim), Δ->(NoTangent(), xsum(Δ), NoTangent())
end

struct NonDiffEven{N, O, P}; end
struct NonDiffOdd{N, O, P}; end

(::NonDiffOdd{N, O, P})(Δ) where {N, O, P} = (ntuple(_->ZeroTangent(), N), NonDiffEven{N, plus1(O), P}())
(::NonDiffEven{N, O, P})(Δ...) where {N, O, P} = (ZeroTangent(), NonDiffOdd{N, plus1(O), P}())
(::NonDiffOdd{N, O, O})(Δ) where {N, O} = ntuple(_->ZeroTangent(), N)

# This should not happen
(::NonDiffEven{N, O, O})(Δ...) where {N, O} = error()

@Base.pure function ChainRulesCore.rrule(::typeof(Core.apply_type), head, args...)
    Core.apply_type(head, args...), NonDiffOdd{plus1(plus1(length(args))), 1, 1}()
end

function ChainRulesCore.rrule(::typeof(Core.tuple), args...)
    Core.tuple(args...), Δ->Core.tuple(NoTangent(), Δ...)
end

# TODO: What to do about these integer rules
@ChainRulesCore.non_differentiable Base.rem(a::Integer, b::Type)

ChainRulesCore.canonicalize(::ChainRulesCore.ZeroTangent) = ChainRulesCore.ZeroTangent()

# Skip AD'ing through the axis computation
function ChainRules.rrule(::typeof(Base.Broadcast.instantiate), bc::Base.Broadcast.Broadcasted)
    return Base.Broadcast.instantiate(bc), Δ->begin
        Core.tuple(NoTangent(), Δ)
    end
end

using StaticArrays

# Force the static arrays constructor to use a vector representation of
# the cotangent space.

struct to_tuple{N}; end
@generated function (::to_tuple{N})(Δ) where {N}
    :( (NoTangent(), Core.tuple( $( ( :(Δ[$i]) for i = 1:N )...) )) )
end
(::to_tuple)(Δ::SArray) = getfield(Δ, :data)

function ChainRules.rrule(::Type{SArray{S, T, N, L}}, x::NTuple{L,T}) where {S, T, N, L}
    SArray{S, T, N, L}(x), to_tuple{L}()
end

function ChainRules.rrule(::Type{SArray{S, T, N, L}}, x::NTuple{L,Any}) where {S, T, N, L}
    SArray{S, T, N, L}(x), to_tuple{L}()
end

function ChainRules.frule((_, ∂x), ::Type{SArray{S, T, N, L}}, x::NTuple{L,T}) where {S, T, N, L}
    SArray{S, T, N, L}(x), SArray{S, T, N, L}(∂x)
end

function ChainRules.frule((_, ∂x), ::Type{SArray{S, T, N, L}}, x::NTuple{L,Any}) where {S, T, N, L}
    SArray{S, T, N, L}(x), SArray{S}(∂x)
end

@ChainRulesCore.non_differentiable StaticArrays.promote_tuple_eltype(T)

function ChainRules.frule((_, ∂A), ::typeof(getindex), A::AbstractArray, args...)
    getindex(A, args...), getindex(∂A, args...)
end

function ChainRules.rrule(::typeof(map), ::typeof(+), A::AbstractArray, B::AbstractArray)
    map(+, A, B), Δ->(NoTangent(), NoTangent(), Δ, Δ)
end

function ChainRules.rrule(::typeof(map), ::typeof(+), A::AbstractVector, B::AbstractVector)
    map(+, A, B), Δ->(NoTangent(), NoTangent(), Δ, Δ)
end

function ChainRules.rrule(AT::Type{<:Array{T,N}}, x::AbstractArray{S,N}) where {T,S,N}
    # We're leaving these in the eltype that the cotangent vector already has.
    # There isn't really a good reason to believe we should convert to the
    # original array type, so don't unless explicitly requested.
    AT(x), Δ->(NoTangent(), Δ)
end

function ChainRules.rrule(AT::Type{<:Array}, undef::UndefInitializer, args...)
    # We're leaving these in the eltype that the cotangent vector already has.
    # There isn't really a good reason to believe we should convert to the
    # original array type, so don't unless explicitly requested.
    AT(undef, args...), Δ->(NoTangent(), NoTangent(), ntuple(_->NoTangent(), length(args))...)
end

function unzip_tuple(t::Tuple)
    map(x->x[1], t), map(x->x[2], t)
end

function ChainRules.rrule(::typeof(unzip_tuple), args::Tuple)
    unzip_tuple(args), Δ->(NoTangent(), map((x,y)->(x,y), Δ...))
end

struct BackMap{T}
    f::T
end
(f::BackMap{N})(args...) where {N} = ∂⃖¹(getfield(f, :f), args...)
back_apply(x, y) = x(y)
back_apply_zero(x) = x(Zero())

function ChainRules.rrule(::typeof(map), f, args::Tuple)
    a, b = unzip_tuple(map(BackMap(f), args))
    function back(Δ)
        (fs, xs) = unzip_tuple(map(back_apply, b, Δ))
        (NoTangent(), sum(fs), xs)
    end
    function back(Δ::ZeroTangent)
        (fs, xs) = unzip_tuple(map(back_apply_zero, b))
        (NoTangent(), sum(fs), xs)
    end
    a, back
end

function ChainRules.rrule(::typeof(Base.ntuple), f, n)
    a, b = unzip_tuple(ntuple(BackMap(f), n))
    a, function (Δ)
        (NoTangent(), sum(map(back_apply, b, Δ)), NoTangent())
    end
end

function ChainRules.frule(_, ::Type{Vector{T}}, undef::UndefInitializer, dims::Int...) where {T}
    Vector{T}(undef, dims...), zeros(T, dims...)
end

@ChainRules.non_differentiable Base.:(|)(a::Integer, b::Integer)
@ChainRules.non_differentiable Base.throw(err)
@ChainRules.non_differentiable Core.Compiler.return_type(args...)
ChainRulesCore.canonicalize(::NoTangent) = NoTangent()

# Disable thunking at higher order (TODO: These should go into ChainRulesCore)
function ChainRulesCore.rrule(::Type{Thunk}, thnk)
    z, ∂z = ∂⃖¹(thnk)
    z, Δ->(NoTangent(), ∂z(Δ)...)
end

function ChainRulesCore.rrule(::Type{InplaceableThunk}, add!!, val)
    val, Δ->(NoTangent(), NoTangent(), Δ)
end


"""

new optimazation idea  - try to put all data in boolean arrays in shared memory  when getting means
next we would need only to read shared memory - yet first one need to check wheather there would be enough shmem on device

calculating intercalss correlation
"""
module InterClassCorrKernel

using CUDA, ..CUDAGpuUtils ,..IterationUtils,..ReductionUtils , ..MemoryUtils,..CUDAAtomicUtils

export prepareInterClassCorrKernel,calculateInterclassCorr

"""
preparation for the interclass correlation kernel - we prepare the  GPU arrays - in this kernel they are small , calculate required loop iterations
and using occupancy API calculate optimal number of  threads and blocks for both kernels
also return arguments lists for both kernels
"""
function prepareInterClassCorrKernel()
  mainArrDims= (2,2,2)
  sumOfGold= CUDA.zeros(1);
  sumOfSegm= CUDA.zeros(1);
  sswTotal= CUDA.zeros(1);
  ssbTotal= CUDA.zeros(1);
  numberToLooFor= 1
  totalNumbOfVoxels= (mainArrDims[1]*mainArrDims[2]*mainArrDims[3])
  pixPerSlice = mainArrDims[1]*mainArrDims[2]
  iterLoop=5
  argsMain=(sumOfGold,sumOfSegm
     ,sswTotal
     ,ssbTotal
     ,iterLoop,pixPerSlice,totalNumbOfVoxels
     ,numberToLooFor )
  get_shmem(threads) = 4*33+3  #the same for both kernels
  threads,blocks =getThreadsAndBlocksNumbForKernel(get_shmem, kernel_InterClassCorr,(CUDA.zeros(2,2,2),CUDA.zeros(2,2,2),argsMain...))
  totalNumbOfVoxels= (mainArrDims[1]*mainArrDims[2]*mainArrDims[3])
  pixPerSlice= cld(totalNumbOfVoxels,blocks)
  iterLoop = UInt32(fld(pixPerSlice, threads[1]*threads[2]))

  argsMain=(sumOfGold,sumOfSegm
     ,sswTotal
     ,ssbTotal
     ,iterLoop,pixPerSlice,totalNumbOfVoxels
     ,numberToLooFor )

  return(argsMain, threads,blocks, totalNumbOfVoxels)
end

"""
calculates slicewise and global interclass correlation metric
"""
function calculateInterclassCorr(flatGold
                                ,flatSegm,threads,blocks
                                ,args,numberToLooFor)::Float64
 mainArrDims= size(flatGold)
 totalNumbOfVoxels= length(flatGold)
 pixPerSlice= cld(totalNumbOfVoxels,blocks)
 iterLoop = UInt32(fld(pixPerSlice, threads[1]*threads[2]))
 #resetting
 for i in  1:4
  CUDA.fill!(args[i],0)
 end
 argsMain=(args[1],args[2]
 ,args[3]
 ,args[4]
 ,iterLoop,pixPerSlice,totalNumbOfVoxels
 ,numberToLooFor )

@cuda threads=threads blocks=blocks cooperative = true kernel_InterClassCorr(flatGold,flatSegm,argsMain...)

  # println("grandMean[1] $(grandMean[1])  \n")
# @cuda threads=threads blocks=blocks  kernel_InterClassCorr(flatGold  ,flatSegm,args... )
     ssw = args[3][1]/args[7];
     ssb = args[4][1]/(args[7]-1) * 2;

  return (ssb - ssw)/(ssb + ssw);

end

"""
adapted from https://github.com/JuliaGPU/CUDA.jl/blob/afe81794038dddbda49639c8c26469496543d831/src/mapreduce.jl
goldBoolGPU - array holding data of gold standard bollean array
segmBoolGPU - boolean array with the data we want to compare with gold standard
ic - holding single value for global interclass correlation
intermediateRes- array holding slice wise results for ic
loopNumb - number of times the single lane needs to loop in order to get all needed data
sliceEdgeLength - length of edge of the slice we need to square this number to get number of pixels in a slice
amountOfWarps - how many warps we can stick in the block
"""

function kernel_InterClassCorr(flatGold  ,flatSegm,sumOfGold,sumOfSegm
       ,sswTotal
       ,ssbTotal
       ,iterLoop,pixPerSlice,totalNumbOfVoxels
       ,numberToLooFor )
   grid_handle = this_grid()

   #for storing results from warp reductions
   shmemSum = @cuStaticSharedMem(Float32, (32,2))   #thread local values that are meant to store some results - like means ...
   grandMeanShmem = @cuStaticSharedMem(Float32, (1))
   offsetIter = UInt8(1)

   locValA = Float32(0)
   locValB = Float32(0)
        #reset shared memory
        @ifY 1 shmemSum[threadIdxX(),1]=0 ;   @ifY 2 shmemSum[threadIdxX(),2]=0
        sync_threads()

        #first we add 1 for each spot we have true - so we will get sum  - and from sum we can get mean
        @iterateLinearlyMultipleBlocks(iterLoop,pixPerSlice,totalNumbOfVoxels,begin
            locValA += flatGold[i]==numberToLooFor
            locValB += flatSegm[i]==numberToLooFor
        end)#for
        @redWitAct(offsetIter,shmemSum,  locValA,+,     locValB,+  )
   sync_threads()
    @addAtomic(shmemSum, sumOfGold,sumOfSegm)
    ### sums should be in place
sync_grid(grid_handle)
grandMeanShmem[1]= ( (sumOfGold[1]/totalNumbOfVoxels) + (sumOfSegm[1]/totalNumbOfVoxels ))/2
locValA = 0
locValB = 0
@ifY 1 shmemSum[threadIdxX(),1]=0 ;
@ifY 2 shmemSum[threadIdxX(),2]=0 ;
sync_threads()

@iterateLinearlyMultipleBlocks(iterLoop,pixPerSlice,totalNumbOfVoxels,begin
m =  ((flatGold[i]==numberToLooFor) +(flatSegm[i]==numberToLooFor))/2
locValA += (((flatGold[i]==numberToLooFor)- m)^2) +(((flatSegm[i]==numberToLooFor)- m )^2)
locValB += ((m - grandMeanShmem[1])^2)
end)#for
#now we accumulated ssw and locValB - we need to reduce it
offsetIter = UInt8(1)
@redWitAct(offsetIter,shmemSum,  locValA,+,     locValB,+  )
sync_threads()
@addAtomic(shmemSum, sswTotal,ssbTotal)
return nothing

end

# function kernel_InterClassCorr(flatGold
#     ,flatSegm,
#     sumOfGold,sumOfSegm
#        ,sswTotal
#        ,ssbTotal
#        ,iterLoop,pixPerSlice,totalNumbOfVoxels
#        ,numberToLooFor,grandMean )

#     #for storing results from warp reductions
#     shmemSum = @cuStaticSharedMem(Float32, (33,2))   #thread local values that are meant to store some results - like means ...
#     @ifY 1 shmemSum[threadIdxX(),1]=0 ;
#     @ifY 2 shmemSum[threadIdxX(),2]=0 ;
#     sync_threads()
#     locValA = Float32(0.0)
#     locValB = Float32(0.0)

#     @iterateLinearlyMultipleBlocks(iterLoop,pixPerSlice,totalNumbOfVoxels,begin
#         m =  ((flatGold[i]==numberToLooFor) +(flatSegm[i]==numberToLooFor))/2
#         locValA += (((flatGold[i]==numberToLooFor)- m)^2) +(((flatSegm[i]==numberToLooFor)- m )^2)
#         locValB += ((m - grandMean[1])^2)
#     end)#for
#     #now we accumulated ssw and locValB - we need to reduce it
#     offsetIter = UInt8(1)
#     @redWitAct(offsetIter,shmemSum,  locValA,+,     locValB,+  )
#    sync_threads()
#     @addAtomic(shmemSum, sswTotal,ssbTotal)
#     return nothing

# end

# function kernel_InterClassCorr_means(goldlGPU,segmGPU
#      ,sumOfGold,sumOfSegm
#      ,sswTotal
#      ,ssbTotal
#      ,iterLoop,pixPerSlice,totalNumbOfVoxels
#      ,numberToLooFor )

# locValA = UInt32(0)
# locValB = UInt32(0)
# offsetIter= UInt8(1)
# shmemSum = @cuStaticSharedMem(UInt16, (33,2))   #thread local values that are meant to store some results - like means ...
# @ifY 1 shmemSum[threadIdxX(),1]= 0
# @ifY 2 shmemSum[threadIdxX(),2]= 0
# sync_threads()

# @iterateLinearlyMultipleBlocks(iterLoop,pixPerSlice,totalNumbOfVoxels,
# #inner expression
# begin
# #updating variables needed to calculate means
# locValA += goldlGPU[i]==numberToLooFor
# locValB += segmGPU[i]==numberToLooFor
# end)

# #tell what variables are to be reduced and by what operation
# @redWitAct(offsetIter,shmemSum,  locValA,+,     locValB,+  )
# sync_threads()
# @addAtomic(shmemSum, sumOfGold, sumOfSegm)

# return nothing
# end

# function kernel_InterClassCorr(goldlGPU,segmGPU
#   ,sumOfGold,sumOfSegm
#   ,sswTotal
#   ,ssbTotal
#   ,iterLoop,pixPerSlice,totalNumbOfVoxels
#   ,numberToLooFor )
# grandMean = @cuStaticSharedMem(Float32, (1))
# @ifXY 1 3 grandMean[1]= ( (sumOfGold[1]/totalNumbOfVoxels) + (sumOfSegm[1]/totalNumbOfVoxels ))/2
# ssw = Float32(0)
# ssb = Float32(0)
# m = Float32(0)
# offsetIter= UInt8(1)
# shmemSum = @cuStaticSharedMem(Float32, (33,2))   #thread local values that are meant to store some results - like means ...
# @ifY 1 shmemSum[threadIdxX(),1]= 0
# @ifY 2 shmemSum[threadIdxX(),2]= 0
# sync_threads()
# @iterateLinearlyMultipleBlocks(iterLoop,pixPerSlice,totalNumbOfVoxels,
#  #inner expression
#  begin
#   m =  ((goldlGPU[i]==numberToLooFor) +(segmGPU[i]==numberToLooFor))/2
#   ssw += (((goldlGPU[i]==numberToLooFor)- m)^2) +(((segmGPU[i]==numberToLooFor)- m )^2)
#   ssb += ((m - grandMean[1])^2)
#  end)

#  #tell what variables are to be reduced and by what operation
#  @redWitAct(offsetIter,shmemSum,  ssw,+,     ssb,+  )
# sync_threads()
# @addAtomic(shmemSum, sswTotal,ssbTotal)

#  return nothing

# end

# """
# calculates slicewise and global interclass correlation metric

# threadsForMeans, threadsForMain - tuples indicating dimensionality  of thread block
# bsically all needed arguments apart from gold standard and algorithm 3 dm arrays  and numberToLooFor should be given by prepareInterClassCorrKernel
# """
# function calculateInterclassCorr(goldGPU,segmGPU,argsMain, threadsMain,  blocksMain,threadsMean,blocksMean,argsMean, totalNumbOfVoxels)::Float64

#   ## resetting some entries to 0
# for i in 1:4
#   CUDA.fill!(argsMain[i],0)
# end

# @cuda threads=threadsMean  blocks=blocksMean  kernel_InterClassCorr_means(goldGPU,segmGPU,argsMean...)
# @cuda threads=threadsMain  blocks=blocksMain  kernel_InterClassCorr(goldGPU,segmGPU,argsMain...)

#   print("before correction ssw $(argsMain[4][1]) ssb $(argsMain[5][1])")

#      ssw = argsMain[4][1]/totalNumbOfVoxels;
#      ssb = (argsMain[5][1]/(totalNumbOfVoxels-1)) * 2;
#     print("ssw $(ssw) ssb $(ssb)")
# #  return (ssb - ssw)/(ssb + ssw);
#  return (ssb - ssw)/(ssb + ssw);

# end

# function kernel_InterClassCorr(goldlGPU,segmGPU
#      ,sumOfGold,sumOfSegm
#      ,sswTotal
#      ,ssbTotal
#      ,iterLoop,pixPerSlice,totalNumbOfVoxels
#      ,numberToLooFor )

# grid_handle = this_grid()

#  locValA = Float32(0)
#  locValB = Float32(0)
#   offsetIter= 1
#   grandMean = @cuStaticSharedMem(Float32, (1))
#   grandMeanSq = @cuStaticSharedMem(Float32, (1))
#   halfMingrandSq = @cuStaticSharedMem(Float32, (1))
#   oneMingrandSq = @cuStaticSharedMem(Float32, (1))

#   shmemSum = @cuStaticSharedMem(Float32, (33,2))   #thread local values that are meant to store some results - like means ...
#   @ifY 1 shmemSum[threadIdxX(),1]= Float32(0)
#   @ifY 2 shmemSum[threadIdxX(),2]= Float32(0)

#   sync_threads()
#   @iterateLinearlyMultipleBlocks(iterLoop,pixPerSlice,totalNumbOfVoxels,
#     #inner expression
#     begin
#         #updating variables needed to calculate means
#            locValA += goldlGPU[i]==numberToLooFor
#            locValB += segmGPU[i]==numberToLooFor
#     end)

#     #tell what variables are to be reduced and by what operation
#     @redWitAct(offsetIter,shmemSum,  locValA,+,     locValB,+  )
#   sync_threads()
#   @addAtomic(shmemSum, sumOfGold, sumOfSegm)

#   # @ifX 1  if(threadIdxY()>22 ) CUDA.@cuprint "aaa idY $(threadIdxY()) \n" end

#  sync_grid(grid_handle)
# #by now we should have all means
#  @ifXY 1 3 grandMean[1]= ( (sumOfGold[1]/totalNumbOfVoxels) + (sumOfSegm[1]/totalNumbOfVoxels ))/2
#  @ifXY 2 3 grandMeanSq[1]= ((( (sumOfGold[1]/totalNumbOfVoxels) + (sumOfSegm[1]/totalNumbOfVoxels ))/2))^2
#  @ifXY 3 3 halfMingrandSq[1]= (0.5- (( (sumOfGold[1]/totalNumbOfVoxels) + (sumOfSegm[1]/totalNumbOfVoxels ))/2))^2
#  @ifXY 4 3 oneMingrandSq[1]= (1- (( (sumOfGold[1]/totalNumbOfVoxels) + (sumOfSegm[1]/totalNumbOfVoxels ))/2))^2

#  @ifXY 1 1 CUDA.@cuprint "  "

#  @ifY 1 shmemSum[threadIdxX(),1]= Float32(0)
#  @ifY 2 shmemSum[threadIdxX(),2]= Float32(0)
#  sync_threads()
#  locValA = 0
#  locValB = 0
#   # m = Float32(0)
# sync_threads()

#   @iterateLinearlyMultipleBlocks(iterLoop,pixPerSlice,totalNumbOfVoxels,
#     begin
#       boolGold = segmGPU[i]==numberToLooFor
#       boolSegm = goldlGPU[i]==numberToLooFor
#       if(boolGold && boolSegm)
#         # CUDA.@cuprint "aa    "
#         locValB = locValB+oneMingrandSq[1]
#       elseif(boolGold ⊻ boolSegm)
#        locValA =locValA+ (0.25)
#        locValB += halfMingrandSq[1]#((0.5 - grandMean[1])*(0.5 - grandMean[1])   )
#       else
#         locValB =locValB+ grandMeanSq[1]
#       end

#       # elseif(boolGold || boolSegm)
#       #   m = 0.5#((boolGold ) +(boolSegm) )/2
#       #   locValA += 0.25
#       #   locValB += ((0.5 - grandMean[1])^2)
#       # else
#       #   locValB += grandMean[1]^2

#       # end
#     end )
#     #tell what variables are to be reduced and by what operation
#       offsetIter= 1
#       @redWitAct(offsetIter,shmemSum,  locValA,+,     locValB,+  )
#    sync_threads()
#     @addAtomic(shmemSum, sswTotal,ssbTotal)
#     #         #now we will have sum of all entries in given slice that comply to our predicate
#     #         #next we need to reduce values
#     #         while(offsetIter <32)
#     #             @inbounds locValA+=shfl_down_sync(FULL_MASK, locValA, offsetIter)
#     #             @inbounds locValB+=shfl_down_sync(FULL_MASK, locValB, offsetIter)
#     #             offsetIter<<= 1
#     #         end
#     #         # @ifX 1  if(threadIdxY()>22 ) CUDA.@cuprint "bbb idY $(threadIdxY()) \n" end
#     #         #@ifX 1  if(blockIdxX()==1 ) CUDA.@cuprint "bbb idY $(threadIdxY()) \n" end
#     #         #@ifY 1  if(blockIdxX()==1 ) CUDA.@cuprint "bbb idX $(threadIdxX()) \n" end

#     #         # @ifX 1  if(blockIdxX()==1 ) CUDA.@cuprint "bbb idY $(threadIdxY()) \n" end
#     #         #@ifY 1  if(blockIdxX()==1 ) CUDA.@cuprint "bbb idX $(threadIdxX()) \n" end

#     #  #   now we have sums in first threads of the warp we need to pass it to shared memory
#     #     if(threadIdxX()==1)
#     #       #TODO() try to do full reduction without those atomics
#     #         @inboundsCUDA.@atomic sswTotal[1]+=locValA
#     #         @inboundsCUDA.@atomic ssbTotal[1]+=locValB
#     #     end
#         # sync_threads()

#         # sync_warp(active_mask())

#         # #finally reduce from shared memory
#         # if(threadIdxY()==1)
#         #     offsetIter = UInt8(1)
#         #     while(offsetIter <32)
#         #         @inbounds shmemSum[threadIdxX(),1]+=shfl_down_sync(FULL_MASK, shmemSum[threadIdxX(),1], offsetIter)
#         #         @inbounds shmemSum[threadIdxX(),2]+=shfl_down_sync(FULL_MASK, shmemSum[threadIdxX(),2], offsetIter)
#         #     offsetIter<<= 1
#         #     end
#         # end

# # @ifXY 1 1 CUDA.@cuprint "shmemSum[1,1]  $(shmemSum[1,1])  \n "
#   #  sync_threads()
#   #  @addAtomic(shmemSum, sswTotal,ssbTotal)

#     return nothing

# end

end#InterClassCorrKernel

#     #offset for lloking for values in source arrays
#     offset = (pixelNumberPerSlice*(blockIdx().x-1))
#     #for storing results from warp reductions
#     shmemSum = @cuStaticSharedMem(Float32, (33,2))   #thread local values that are meant to store some results - like means ...
#     @ifY 1 shmemSum[threadIdxX(),1]=0 ;   @ifY 2 shmemSum[threadIdxX(),2]=0 ;@ifY 3 shmemSum[threadIdxX(),2]=0
#     sync_threads()
#     ssw::Float32 = Float32(0.0)
#     ssb::Float32 = Float32(0.0)

#     @unroll for k in UInt16(0):loopNumb
#         if(threadIdxX()+(threadIdxY()-1)*32+k*1024 <=pixelNumberPerSlice)
#         ind =offset+ threadIdxX()+(threadIdxY()-1)*32+k*1024
#         m =  ((flatGold[ind]==numberToLooFor) +(flatSegm[ind]==numberToLooFor))/2
#         ssw += (((flatGold[ind]==numberToLooFor)- m)^2) +(((flatSegm[ind]==numberToLooFor)- m )^2)
#         ssb += ((m - grandMean[1])^2)
#         end#if
#     end#for
#     #now we accumulated ssw and ssb - we need to reduce it
#     offsetIter = UInt8(1)
#     while(offsetIter <32)
#         @inbounds ssw+=shfl_down_sync(FULL_MASK, ssw, offsetIter)
#         @inbounds ssb+=shfl_down_sync(FULL_MASK, ssb, offsetIter)
#         offsetIter<<= 1
#     end
#         # now we have sums in first threads of the warp we need to pass it to shared memory
#         if(threadIdxX()==1)
#             @inbounds shmemSum[threadIdxY(),1]+=ssw
#             @inbounds shmemSum[threadIdxY(),2]+=ssb
#         end
#     sync_threads()
#     #finally reduce from shared memory
#         if(threadIdxY()==1)
#             offsetIter = UInt8(1)
#             while(offsetIter <32)
#             @inbounds shmemSum[threadIdxX(),1]+=shfl_down_sync(FULL_MASK, shmemSum[threadIdxX(),1], offsetIter)
#             @inbounds shmemSum[threadIdxX(),2]+=shfl_down_sync(FULL_MASK, shmemSum[threadIdxX(),2], offsetIter)
#             offsetIter<<= 1
#             end
#     end
#     sync_threads()sswTotal,ssbTotal
      #now in   shmemSum[1,1] we should have ssw and in  shmemSum[1,2] ssb
#       @ifXY 1 1  @inboundsCUDA.@atomic sswTotal[]+= shmemSum[1,1]
#       @ifXY 1 2  @inboundsCUDA.@atomic ssbTotal[]+= shmemSum[1,2]
#       @ifXY 1 3  begin
#         sswInner = shmemSum[1,2]/totalNumbOfVoxels;
#         ssbInner = shmemSum[1,1]/(totalNumbOfVoxels-1) * 2
#         @inbounds iccPerSlice[blockIdxX()]=(ssb - ssw)/(ssb + ssw)
#       end
    # # ####### now we have ssw and ssb calculated both global and per slice

# """
# adapted from https://github.com/JuliaGPU/CUDA.jl/blob/afe81794038dddbda49639c8c26469496543d831/src/mapreduce.jl
# goldGPU - array holding data of gold standard  array
# segmGPU -  array with the data we want to compare with gold standard
# ic - holding single value for global interclass correlation
# intermediateRes- array holding slice wise results for ic
# loopXdim,loopYdim,loopZdim - number of times we need to loop over each dimension in order to get all needed data
# sliceEdgeLength - length of edge of the slice we need to square this number to get number of pixels in a slice
# amountOfWarps - how many warps we can stick in the block
# """

# function kernel_InterClassCorr_means(goldlGPU
#                                ,segmGPU
#                                 ,iterLoop,pixPerSlice,totalNumbOfVoxels
#                                 ,sumOfGold
#                                 ,sumOfSegm
#                                 ,numberToLooFor )

#   locValA = UInt32(0)
#  locValB = UInt32(0)
#   offsetIter= UInt8(1)
#   shmemSum = @cuStaticSharedMem(UInt16, (33,2))   #thread local values that are meant to store some results - like means ...
#   @iterateLinearlyMultipleBlocks(iterLoop,pixPerSlice,totalNumbOfVoxels,
#     #inner expression
#     begin
#         #updating variables needed to calculate means
#            locValA += goldlGPU[i]==numberToLooFor
#            locValB += segmGPU[i]==numberToLooFor
#     end)

#     #tell what variables are to be reduced and by what operation
#     @redWitAct(offsetIter,shmemSum,  locValA,+,     locValB,+  )
#   sync_threads()
#   @addAtomic(shmemSum, sumOfGold, sumOfSegm)

# #     #offset for lloking for values in source arrays
# #     offset = (pixelNumberPerSlice*(blockIdx().x-1))
# #    #for storing results from warp reductions
# #    shmemSum = @cuStaticSharedMem(UInt16, (33,2))   #thread local values that are meant to store some results - like means ...
# #    offsetIter = UInt8(1)

# #    locValA = UInt32(0)
# #    locValB = UInt32(0)
# #         #reset shared memory
# #         @ifY 1 shmemSum[threadIdxX(),1]=0 ;   @ifY 2 shmemSum[threadIdxX(),2]=0
# #         sync_threads()

# #         #first we add 1 for each spot we have true - so we will get sum  - and from sum we can get mean
# #         @unroll for k in 0:loopNumb
# #             if(threadIdxX()+(threadIdxY()-1)*32+k*1024 <=pixelNumberPerSlice)
# #             ind =offset+ threadIdxX()+(threadIdxY()-1)*32+k*1024
# #             locValA += flatGold[ind]==numberToLooFor
# #             locValB += flatSegm[ind]==numberToLooFor
# #             end#if
# #         end#for

# #             #now we will have sum of all entries in given slice that comply to our predicate
# #             #next we need to reduce values
# #             offsetIter = UInt8(1)
# #             while(offsetIter <32)
# #                 @inbounds locValA+=shfl_down_sync(FULL_MASK, locValA, offsetIter)
# #                 @inbounds locValB+=shfl_down_sync(FULL_MASK, locValB, offsetIter)
# #                 offsetIter<<= 1
# #             end
# #         # now we have sums in first threads of the warp we need to pass it to shared memory
# #         if(threadIdxX()==1)

# #             @inbounds shmemSum[threadIdxY(),1]+=locValA
# #             @inbounds shmemSum[threadIdxY(),2]+=locValB
# #         end
# #         sync_threads()
# #         #finally reduce from shared memory
# #         if(threadIdxY()==1)
# #             offsetIter = UInt8(1)
# #             while(offsetIter <32)
# #                 @inbounds shmemSum[threadIdxX(),1]+=shfl_down_sync(FULL_MASK, shmemSum[threadIdxX(),1], offsetIter)
# #                 @inbounds shmemSum[threadIdxX(),2]+=shfl_down_sync(FULL_MASK, shmemSum[threadIdxX(),2], offsetIter)
# #             offsetIter<<= 1
# #             end
# #         end
# # sync_threads()

# #       #now in   shmemSum[1,1] we should have sum of values complying with our predicate in gold mask and in shmemSum[1,2] values of other mask
# #       #we need to add now those to the globals
# #       @ifXY 1 1  @inboundsCUDA.@atomic sumOfGold[]+= shmemSum[1,1]
# #       @ifXY 1 2  @inboundsCUDA.@atomic sumOfSegm[]+= shmemSum[1,2]

#     return nothing
# end


"""
    energy_balance(object::T,meteo::Atmosphere,constants = Constants())
    energy_balance(object::Array{AbstractComponentModel},meteo::Atmosphere,constants = Constants())
    energy_balance!(object::AbstractComponentModel,meteo::Atmosphere,constants = Constants())

Computes the energy balance of a component based on the type of the model it was parameterized
with in `object.energy`.

At the moment, two models are implemented in the package:

- [`Monteith`](@ref): the model found in Monteith and Unsworth (2013)
- `Missing`: if no computation of the energy balance is needed

# Arguments

- `object::Union{AbstractComponentModel,Array{AbstractComponentModel},
Dict{String,AbstractComponentModel}}`: a [`Component`](@ref) struct, or a Dict/Array of.
- `meteo::Union{Atmosphere,Weather}`: meteorology structure, see [`Atmosphere`](@ref) or
[`Weather`](@ref)
- `constants = Constants()`: physical constants. See [`Constants`](@ref) for more details

# Note

Some models need input values for some variables. For example [`Monteith`](@ref) requires a
value for `Rₛ`, `d` and `skyFraction`. If you read the models from a file, you can
use [`init_status!`](@ref) (see examples).

# Examples

```julia
meteo = Atmosphere(T = 20.0, Wind = 1.0, P = 101.3, Rh = 0.65)

# Using the model of Monteith and Unsworth (2013) for energy, Farquhar et al. (1980) for
# photosynthesis, and Medlyn et al. (2011) for stomatal conductance:
leaf = LeafModels(energy = Monteith(),
            photosynthesis = Fvcb(),
            stomatal_conductance = Medlyn(0.03, 12.0),
            Rₛ = 13.747, skyFraction = 1.0, PPFD = 1500.0, d = 0.03)

energy_balance(leaf,meteo)

# Using a model file:
model = read_model("a-model-file.yml")

# An example model file is available here:
# "https://raw.githubusercontent.com/VEZY/PlantBiophysics/main/test/inputs/models/plant_coffee.yml"

# Initialising the mandatory variables:
init_status!(model, Rₛ = 13.747, skyFraction = 1.0, PPFD = 1500.0, Tₗ = 25.0, d = 0.03)

# NB: To know which variables has to be initialised according to the models used, you can use
# `to_initialise(leaf)`

# Running a simulation for all component types in the same scene:
energy_balance!(model, meteo)

model["Leaf"].status.Rn
model["Leaf"].status.A
model["Leaf"].status.Cᵢ
```

# References

Duursma, R. A., et B. E. Medlyn. 2012. « MAESPA: a model to study interactions between water
limitation, environmental drivers and vegetation function at tree and stand levels, with an
example application to [CO2] × drought interactions ». Geoscientific Model Development 5 (4):
919‑40. https://doi.org/10.5194/gmd-5-919-2012.

Monteith, John L., et Mike H. Unsworth. 2013. « Chapter 13 - Steady-State Heat Balance: (i)
Water Surfaces, Soil, and Vegetation ». In Principles of Environmental Physics (Fourth Edition),
edited by John L. Monteith et Mike H. Unsworth, 217‑47. Boston: Academic Press.

Schymanski, Stanislaus J., et Dani Or. 2017. « Leaf-Scale Experiments Reveal an Important
Omission in the Penman–Monteith Equation ». Hydrology and Earth System Sciences 21 (2): 685‑706.
https://doi.org/10.5194/hess-21-685-2017.

Vezy, Rémi, Mathias Christina, Olivier Roupsard, Yann Nouvellon, Remko Duursma, Belinda Medlyn,
Maxime Soma, et al. 2018. « Measuring and modelling energy partitioning in canopies of varying
complexity using MAESPA model ». Agricultural and Forest Meteorology 253‑254 (printemps): 203‑17.
https://doi.org/10.1016/j.agrformet.2018.02.005.
"""
function energy_balance!(object::AbstractComponentModel, meteo::Atmosphere, constants = Constants())
    is_init = is_initialised(object)
    !is_init && error("Some variables must be initialized before simulation")

    net_radiation!(object, meteo, constants)
    return nothing
end

# Same as above but non-mutating
function energy_balance(object::AbstractComponentModel, meteo::Atmosphere, constants = Constants())
    object_tmp = copy(object)
    energy_balance!(object_tmp, meteo, constants)
    return object_tmp.status
end

# energy_balance over several objects (e.g. all leaves of a plant) in an Array
function energy_balance!(object::O, meteo::Atmosphere, constants = Constants()) where O <: AbstractArray{<:AbstractComponentModel}

    for i in values(object)
        energy_balance!(i, meteo, constants)
    end

    return nothing
end

# same as the above but non-mutating
function energy_balance(
    object::O,
    meteo::M,
    constants = Constants()
    ) where {O <: Union{AbstractArray{<:AbstractComponentModel},AbstractDict{N,<:AbstractComponentModel} where N},M <: Union{Atmosphere,Weather}}

    # Copy the objects only once before the computation for performance reasons:
    object_tmp = copy(object)

    # Computation:
    energy_balance!(object_tmp, meteo, constants)

    # --- Extracting the outputs: ---

    # Pre-allocating the outputs:
    output = [i.status for i in object_tmp]

    for (i, obj) in enumerate(object_tmp)
        output[i] = obj.status
    end

    output = DataFrame([NamedTuple(i) for i in output])

    return output
end

# energy_balance over several objects (e.g. all leaves of a plant) in a kind of Dict.
function energy_balance!(object::O, meteo::Atmosphere, constants = Constants()) where O <: AbstractDict{N,<:AbstractComponentModel} where N

    for (k, v) in object
        energy_balance!(v, meteo, constants)
    end

    return nothing
end

# same as the above but non-mutating. In this case we add a column with the component name 😃
function energy_balance(
    object::O,
    meteo::M,
    constants = Constants()
    ) where {O <: Union{AbstractArray{<:AbstractComponentModel},AbstractDict{N,<:AbstractComponentModel} where N},M <: Union{Atmosphere,Weather}}

    # Copy the objects only once before the computation for performance reasons:
    object_tmp = copy(object)

    # Computation:
    energy_balance!(object_tmp, meteo, constants)

    # --- Extracting the outputs: ---

    # Pre-allocating the outputs:
    output = Dict([k => v.status for (k, v) in object_tmp])

    for (k, v) in object_tmp
        output[k] = v.status
    end

    output = DataFrame([(NamedTuple(v)..., component = k) for (k, v) in output])

    return output
end

# energy_balance over several meteo time steps (called Weather) and possibly several components.
# Only allowed for components given as a subtype of AbstractDict to track components names in the
# outputs
function energy_balance!(
    object::T,
    meteo::Weather,
    constants = Constants()
    ) where T <: Union{AbstractDict{N,<:AbstractComponentModel} where N}

    # Pre-allocating the time-step outputs:
    timestep_tmp = Dict([k => v.status for (k, v) in object])

    # Pre-allocating the general DataFrame with the first time-step results:
    energy_balance!(object, meteo.data[1], constants)

    for (k, v) in object
        timestep_tmp[k] = v.status
    end

    output_timestep = DataFrame([(NamedTuple(v)..., component = k) for (k, v) in timestep_tmp])

    # Actually pre-allocating the DF:
    output = repeat(output_timestep, length(meteo.data))
    output.time_step = repeat(1:length(meteo.data), inner = size(output_timestep, 1))

    # Computing for all following time-steps:
    for (i, meteo_i) in enumerate(meteo.data[2:end])
        energy_balance!(object, meteo_i, constants)

        for (k, v) in object
            timestep_tmp[k] = v.status
        end

        output_timestep = DataFrame([(NamedTuple(v)..., component = k) for (k, v) in timestep_tmp])

        # Update the values of the global output:
        output[output.time_step .== i,Not(:time_step)] = output_timestep
    end

    return output
end

# energy_balance over several meteo time steps (same as above) but non-mutating
function energy_balance(
    object::T,
    meteo::Weather,
    constants = Constants()
    ) where T <: Union{AbstractDict{N,<:AbstractComponentModel} where N}

    object_tmp = copy(object)

    return energy_balance!(object_tmp, meteo, constants)
end