text:
using Statistics
import Flux

using ..Transformers.Basic: gather, OneHot, OneHotArray

import Flux.Losses

Base.getindex(a::AbstractVector, o::OneHot{K}) where K = a[Int(o)]
Base.getindex(a::AbstractArray, o::OneHot{K}, i...) where K = a[Int(o), i...]

function Base.getindex(xs::AbstractArray{T}, onehots::OneHotArray{K}) where {T, K}
  @assert size(xs, 1) == K
  ys = fill!(similar(xs, Base.tail(size(xs))), zero(T))

  Threads.@threads for i = 1:length(ys)
    @inbounds if (o = Int(onehots.data[i])) != 0
      ind = Tuple(CartesianIndices(onehots.data)[i])
      ys[ind...] = xs[o, ind...]
    end
  end
  return ys
end

using Flux: @adjoint

@adjoint function Base.getindex(xs::AbstractArray, onehots::OneHotArray{K}) where K
  xs[onehots], Δ -> begin
    dxs = fill!(similar(xs), zero(eltype(xs)))
    (∇getindex!(dxs, Δ, onehots), nothing)
  end
end

function ∇getindex!(xs, ys, onehots::OneHotArray{K}) where K
  Threads.@threads for i = 1:length(ys)
    @inbounds if (o = Int(onehots.data[i])) != 0
      ind = Tuple(CartesianIndices(onehots.data)[i])
      xs[o, ind...] = ys[ind...]
    end
  end
  return xs
end

import CUDA
import CUDA: CuVector, CuArray, CuDeviceArray, @cuda

function Base.getindex(xs::CuArray{T}, onehots::OneHotArray{K, N, N2, <:CuArray{OneHot{K}}}) where {T, K, N, N2}
  @assert size(xs, 1) == K
  ys = CUDA.zeros(T, size(onehots.data))

  function kernel(ys::CuDeviceArray{T}, xs::CuDeviceArray{T}, onehots)
    i = CUDA.threadIdx().x + (CUDA.blockIdx().x - 1) * CUDA.blockDim().x

    @inbounds if i <= length(onehots) && (o = Int(onehots[i])) != 0
      ind = Tuple(CartesianIndices(onehots)[i])
      ys[ind...] = xs[o, ind...]
    end
    return
  end

  function configurator(kernel)
    config = launch_configuration(kernel.fun)

    threads = Base.min(length(onehots.data), config.threads)
    blocks = cld(length(onehots.data), threads)

    return (threads=threads, blocks=blocks)
  end

  kernel = @cuda name="onehot_getindex" launch=false kernel(ys, xs, onehots.data)
  config = configurator(kernel)
  kernel(ys, xs, onehots.data; config...)
  return ys
end

function ∇getindex!(xs::CuArray{T}, ys::CuArray{T}, onehots::OneHotArray{K, N, N2, <:CuArray{OneHot{K}}}) where {T, K, N, N2}
  function kernel(xs::CuDeviceArray{T}, ys::CuDeviceArray{T}, onehots)
    i = CUDA.threadIdx().x + (CUDA.blockIdx().x - 1) * CUDA.blockDim().x

    @inbounds if i <= length(onehots) && (o = Int(onehots[i])) != 0
      ind = Tuple(CartesianIndices(onehots)[i])
      xs[o, ind...] = ys[ind...]
    end
    return
  end

  function configurator(kernel)
    config = launch_configuration(kernel.fun)

    threads = Base.min(length(onehots.data), config.threads)
    blocks = cld(length(onehots.data), threads)

    return (threads=threads, blocks=blocks)
  end

  kernel = @cuda name="onehot_getindex_grad" launch=false kernel(xs, ys, onehots.data)
  config = configurator(kernel)
  kernel(xs, ys, onehots.data; config...)
  return xs
end

struct MaskFillArray{T, N, S<:NTuple{N, Int},
                   L<:AbstractVector{T},
                   I<:NTuple{N, <:Integer}, P<:AbstractVector{I}} <: AbstractArray{T, N}
  size::S
  masked_val::T
  values::L
  positions::P
end

MaskFillArray(size, masked_val::T) where T = MaskFillArray(size, masked_val, Vector{T}(), Vector{NTuple{length(size), Int}}())

Base.size(ma::MaskFillArray) = ma.size

function Base.getindex(ma::MaskFillArray{T, N}, i::Vararg{Int, N}) where {T, N}
  @boundscheck checkbounds(ma, i...)
  idx = findfirst(isequal(i), ma.positions)
  if isnothing(idx)
    return ma.masked_val
  else
    return ma.values[idx]
  end
end

function Base.getindex(xs::CuArray{T}, indices::CuVector{<:Integer}) where T
  ys = CuArray{T}(undef, length(indices))

  function kernel(ys::CuDeviceArray{T}, xs, indices)
    i = CUDA.threadIdx().x + (CUDA.blockIdx().x - 1) * CUDA.blockDim().x

    @inbounds if i <= length(indices)
      idx = indices[i]
      ys[i] = xs[idx]
    end
    return
  end

  function configurator(kernel)
    config = launch_configuration(kernel.fun)

    threads = Base.min(length(indices), config.threads)
    blocks = cld(length(indices), threads)

    return (threads=threads, blocks=blocks)
  end

  kernel = @cuda name="linear_getindex" launch=false kernel(ys, xs, indices)
  config = configurator(kernel)
  kernel(ys, xs, indices; config...)
  return ys
end

function Base.getindex(ma::MaskFillArray{T, N}, indices::Vararg{Union{Colon, UnitRange, StepRange, Int}, N}) where {T, N}
  @boundscheck checkbounds(ma, indices...)

  indices_region = map(size(ma), indices) do si, ii
    if ii isa Int
      0
    elseif ii isa Colon
      si
    else
      length(ii)
    end
  end
  remaining_dim = map(!isequal(0), indices_region)
  new_size = filter(!iszero, indices_region)

  function transform_idx(old_idx)
    new_idx = map(old_idx, indices) do idx, ii
      if ii isa Int
        0
      elseif ii isa Colon
        idx
      else
        findfirst(isequal(idx), ii)
      end
    end
    filter(!iszero, new_idx)
  end

  _match(x::Tuple{Colon, Int}) = true
  _match(x::Tuple{Int, Int}) = x[1] == x[2]
  _match(x) = x[2] in x[1]
  match(idx) = all(_match, zip(indices, idx))

  if ma.positions isa CUDA.CuArray
    CUDA.@allowscalar idx = findall(match, ma.positions)
  else
    idx = findall(match, ma.positions)
  end

  new_values = ma.values[idx]
  new_positions = ma.positions[idx]
  if ma.positions isa CUDA.CuArray
    CUDA.@allowscalar new_positions = map(transform_idx, new_positions)
  else
    new_positions = map(transform_idx, new_positions)
  end

  MaskFillArray(new_size, ma.masked_val, new_values, new_positions)
end

function Base.setindex!(ma::MaskFillArray{T, N}, v, i::Vararg{Int, N}) where {T, N}
  @boundscheck checkbounds(ma, i...)
  idx = findfirst(isequal(i), ma.positions)
  if isnothing(idx)
    push!(ma.values, v)
    push!(ma.positions, i)
    return v
  else
    ma.values[idx] = v
  end
end

function MaskFillArray(oma::MaskFillArray{OneHot{K}}) where K
  new_size = (K, size(oma)...)
  new_values = map(x->true, oma.values)
  new_positions = map(zip(oma.values, oma.positions)) do (val, pos)
    (convert(Int, val), pos...)
  end
  MaskFillArray(new_size, false, new_values, new_positions)
end

Flux.@nograd MaskFillArray

import Adapt: adapt, adapt_structure
adapt_structure(T, ma::MaskFillArray) = MaskFillArray(ma.size, ma.masked_val, adapt(T, ma.values), adapt(T, ma.positions))

num_effective_entry(y::OneHotArray) = sum(map(!iszero, y.data))
Flux.@nograd num_effective_entry

function Losses.crossentropy(ŷ::AbstractArray{T, N}, y::OneHotArray{K, N}; agg=mean, ϵ=Flux.epseltype(ŷ)) where {T, K, N}
  ce = (.-log.(ŷ .+ ϵ))[y]
  if agg === sum || agg === mean
    ce = sum(ce)
    if agg === mean
      n = convert(T, num_effective_entry(y))
      return ce / n
    else
      return ce
    end
  else
    return agg(ce)
  end
end

function Losses.logitcrossentropy(ŷ::AbstractArray{T, N}, y::OneHotArray{K, N}; agg=mean) where {T, N, K}
  ce = .-logsoftmax(ŷ)[y]
  if agg === sum || agg === mean
    ce = sum(ce)
    if agg === mean
      n = convert(T, num_effective_entry(y))
      return ce / n
    else
      return ce
    end
  else
    return agg(ce)
  end
end

function Losses.crossentropy(ŷ, y::MaskFillArray{OneHot{K}}; agg=mean, ϵ=Flux.epseltype(ŷ)) where K
  real_ŷ = gather(ŷ, y.positions)
  real_y = OneHotArray(y.values)
  Losses.crossentropy(real_ŷ, real_y; agg=agg)
end

function Losses.logitcrossentropy(ŷ, y::MaskFillArray{OneHot{K}}; agg=mean) where K
  real_ŷ = gather(ŷ, y.positions)
  real_y = OneHotArray(y.values)
  Losses.logitcrossentropy(real_ŷ, real_y; agg=agg)
end


text:
## =========================================================================================
## Initialization Functions -- 0
## =========================================================================================

"""
    $SIGNATURES

Init and start state machine.
"""
function initStartCliqStateMachine!(dfg::AbstractDFG,
                                    tree::AbstractBayesTree,
                                    cliq::TreeClique,
                                    timeout::Union{Nothing, <:Real}=nothing;
                                    oldcliqdata::BayesTreeNodeData=BayesTreeNodeData(),
                                    verbose::Bool=false,
                                    verbosefid=stdout,
                                    drawtree::Bool=false,
                                    show::Bool=false,
                                    incremental::Bool=true,
                                    limititers::Int=20,
                                    upsolve::Bool=true,
                                    downsolve::Bool=true,
                                    recordhistory::Bool=false,
                                    delay::Bool=false,
                                    logger::SimpleLogger=SimpleLogger(Base.stdout),
                                    solve_progressbar=nothing,
                                    algorithm::Symbol=:default,
                                    solveKey::Symbol=algorithm )

  # NOTE use tree and messages for operations involving children and parents
  # TODO deprecate children and prnt clique copies
  # children = TreeClique[]
  # prnt = TreeClique[]

  destType = dfg isa InMemoryDFGTypes ? typeof(dfg) : InMemDFGType

  csmc = CliqStateMachineContainer( dfg, initfg(destType, solverParams=getSolverParams(dfg)),
                                    tree, cliq,
                                    incremental, drawtree, downsolve, delay,
                                    getSolverParams(dfg), Dict{Symbol,String}(), oldcliqdata, logger,
                                    cliq.id, algorithm, 0, true, solveKey, 0)

  !upsolve && !downsolve && error("must attempt either up or down solve")
  # nxt = buildCliqSubgraph_StateMachine
  nxt = setCliqueRecycling_StateMachine

  csmiter_cb = getSolverParams(dfg).drawCSMIters ? ((st::StateMachine)->(cliq.attributes["xlabel"]=st.iter; csmc._csm_iter=st.iter)) : ((st)->(csmc._csm_iter=st.iter))

  statemachine = StateMachine{CliqStateMachineContainer}(next=nxt, name="cliq$(getId(cliq))")

  # store statemachine and csmc in task
  if dfg.solverParams.dbg || recordhistory
    task_local_storage(:statemachine, statemachine)
    task_local_storage(:csmc, csmc)
  end

  logCSM(csmc, "Clique $(getId(csmc.cliq)) starting", loglevel=Logging.Debug)

  #TODO
  # timeout
  # verbosefid=verbosefid
  # injectDelayBefore=injectDelayBefore

  while statemachine(csmc, timeout; verbose=verbose, verbosefid=verbosefid, verboseXtra=getCliqueStatus(csmc.cliq), iterlimit=limititers, recordhistory=recordhistory, housekeeping_cb=csmiter_cb)
    !isnothing(solve_progressbar) && next!(solve_progressbar)
  end

  return CSMHistoryTuple.(statemachine.history)

end

"""
    $SIGNATURES

Recycle clique setup for later uses

Notes
- State machine function 0a
"""
function setCliqueRecycling_StateMachine(csmc::CliqStateMachineContainer)

  oldstatus = getCliqueStatus(csmc.oldcliqdata)

  # canCliqMargRecycle
  if areCliqVariablesAllMarginalized(csmc.dfg, csmc.cliq)
    getCliqueData(csmc.cliq).allmarginalized = true
    setCliqueStatus!(csmc.cliq, MARGINALIZED)

  # canCliqIncrRecycle
  # check if should be trying and can recycle clique computations
  elseif csmc.incremental && oldstatus == DOWNSOLVED
    csmc.cliq.data.isCliqReused = true
    setCliqueStatus!(csmc.cliq, UPRECYCLED)
  end
  logCSM(csmc, "CSM-0a Recycling clique $(csmc.cliqId) from $oldstatus";
              incremental=csmc.cliq.data.isCliqReused,
              marginalized=getCliqueData(csmc.cliq).allmarginalized)

  return buildCliqSubgraph_StateMachine

end

"""
    $SIGNATURES

Build a sub factor graph for clique variables from the larger factor graph.

Notes
- State machine function 0b
"""
function buildCliqSubgraph_StateMachine(csmc::CliqStateMachineContainer)
  # build a local subgraph for inference operations
  syms = getCliqAllVarIds(csmc.cliq)

  logCSM(csmc, "CSM-0b build subgraph syms=$(syms)")

  frontsyms = getCliqFrontalVarIds(csmc.cliq)
  sepsyms = getCliqSeparatorVarIds(csmc.cliq)

  # TODO optimize by only fetching csmc.solveKey -- upgrades required
  buildCliqSubgraph!(csmc.cliqSubFg, csmc.dfg, frontsyms, sepsyms)

  # store the cliqSubFg for later debugging
  _dbgCSMSaveSubFG(csmc, "fg_build")

  # go to 2 wait for up
  return presolveChecklist_StateMachine
end

"""
    $SIGNATURES

Check that the csmc container has everything it needs to proceed with init-ference.

DevNotes
- TODO marginalized flag might be wrong default.
"""
function presolveChecklist_StateMachine(csmc::CliqStateMachineContainer)

  # check if solveKey is available in all variables?
  for var in getVariable.(csmc.cliqSubFg, ls(csmc.cliqSubFg))
    if !(csmc.solveKey in listSolveKeys(var))
      logCSM(csmc, "CSM-0b create empty data for $(getLabel(var)) on solveKey=$(csmc.solveKey)")
      varType = getVariableType(var)
      # FIXME check the marginalization requirements
      setDefaultNodeData!(var, 0, getSolverParams(csmc.cliqSubFg).N, getDimension(varType), solveKey=csmc.solveKey,
                          initialized=false, varType=varType, dontmargin=false)
      #
      @info "create vnd solveKey" csmc.solveKey N
      @info "also" listSolveKeys(var)
    end
  end

  # go to 2 wait for up
  return waitForUp_StateMachine
end

## =========================================================================================
## Wait for up -- 1
## =========================================================================================

"""
    $SIGNATURES

Branching up state
Notes
- State machine function 1
- Common state for handeling messages with take! approach
"""
function waitForUp_StateMachine(csmc::CliqStateMachineContainer)

  logCSM(csmc, "CSM-1 Wait for up messages if needed")

  # setCliqueDrawColor!(csmc.cliq, "olive") #TODO don't know if this is correct color

  # JT empty upRx buffer to save messages, TODO It may be ok not to empty
  beliefMessages = empty!(getMessageBuffer(csmc.cliq).upRx)

  # take! messages from edges
  @sync for e in getEdgesChildren(csmc.tree, csmc.cliq)
    @async begin
      thisEdge = e.dst
      logCSM(csmc, "CSM-1 $(csmc.cliq.id): take! on edge $thisEdge")
      # Blocks until data is available. -- take! model
      beliefMsg = takeBeliefMessageUp!(csmc.tree, e)
      beliefMessages[thisEdge] = beliefMsg
      logCSM(csmc, "CSM-1 $(csmc.cliq.id): Belief message received with status $(beliefMsg.status)"; msgvars = keys(beliefMsg.belief))
    end
  end

  # get all statuses from messages
  all_child_status = map(msg -> msg.status, values(beliefMessages))

  # Main Branching happens here - all up messages received

  # If one up error is received propagate ERROR_STATUS
  if ERROR_STATUS in all_child_status

    putErrorUp(csmc)
    #if its a root, propagate error down
    #FIXME rather check if no parents with function (hasParents or isRoot)
    if length(getParent(csmc.tree, csmc.cliq)) == 0
      putErrorDown(csmc)
      return IncrementalInference.exitStateMachine
    end

    return waitForDown_StateMachine

  elseif csmc.algorithm == :parametric
    !all(all_child_status .== UPSOLVED) && error("#FIXME")
    return solveUp_ParametricStateMachine

  elseif true #TODO Currently all up goes through solveUp
    return preUpSolve_StateMachine

  else
    error("CSM-1 waitForUp State Error: Unknown transision.")
  end

end

## =========================================================================================
## Up functions -- 2
## =========================================================================================

"""
    $SIGNATURES

Notes
- State machine function 2a
"""
function preUpSolve_StateMachine(csmc::CliqStateMachineContainer)

  all_child_status = map(msg -> msg.status, values(getMessageBuffer(csmc.cliq).upRx))

  logCSM(csmc, "CSM-2a preUpSolve_StateMachine with child status"; all_child_status=all_child_status)

  #TODO perhaps don't add for MARGINALIZED
  # always add messages in case its needed for downsolve (needed for differential)
  # add message factors from upRx: cached messages taken from children saved in this clique
  addMsgFactors!(csmc.cliqSubFg, getMessageBuffer(csmc.cliq).upRx, UpwardPass)
  logCSM(csmc, "CSM-2a messages for up"; upmsg=lsf(csmc.cliqSubFg, tags=[:__LIKELIHOODMESSAGE__]))

  # store the cliqSubFg for later debugging
  _dbgCSMSaveSubFG(csmc, "fg_beforeupsolve")

  all_child_finished_up = all(in.(all_child_status, Ref([UPSOLVED, UPRECYCLED, MARGINALIZED])))
  logCSM(csmc, "CSM-2a, clique $(csmc.cliqId) all_child_finished_up $(all_child_finished_up)")

  #try to skip upsolve
  if !getSolverParams(csmc.dfg).upsolve
    return tryDownSolveOnly_StateMachine
  end

  #Clique and children UPSOLVED, UPRECYCLED or MARGINALIZED (finished upsolve)
  #no need to solve
  if getCliqueStatus(csmc.cliq) in [UPSOLVED, UPRECYCLED, MARGINALIZED] && all_child_finished_up
    logCSM(csmc, "CSM-2a Reusing clique $(csmc.cliqId) as $(getCliqueStatus(csmc.cliq))")
    getCliqueStatus(csmc.cliq) == MARGINALIZED &&  setCliqueDrawColor!(csmc.cliq, "blue")
    getCliqueStatus(csmc.cliq) == UPRECYCLED &&  setCliqueDrawColor!(csmc.cliq, "orange")
    return postUpSolve_StateMachine
  end

  # if all(all_child_status .== UPSOLVED)
  if all_child_finished_up
    return solveUp_StateMachine
  elseif !areCliqVariablesAllInitialized(csmc.cliqSubFg, csmc.cliq, csmc.solveKey)
    return initUp_StateMachine
  else
    setCliqueDrawColor!(csmc.cliq, "brown")
    logCSM(csmc, "CSM-2a Clique $(csmc.cliqId) is initialized but children need to init, don't do anything")
    setCliqueStatus!(csmc.cliq, INITIALIZED)
    return postUpSolve_StateMachine
  end
end

"""
  $SIGNATURES

Notes
- State machine function 2b
"""
function initUp_StateMachine(csmc::CliqStateMachineContainer)

  csmc.init_iter += 1

        # FIXME experimental init to whatever is in frontals
        # should work if linear manifold
        # hardcoded off
        linear_on_manifold = false
        init_for_differential = begin
          allvars = getVariables(csmc.cliqSubFg)
          any_init = any(isInitialized.(allvars, csmc.solveKey))
          is_root = isempty(getEdgesParent(csmc.tree, csmc.cliq))
          logCSM(csmc, "CSM-2b init_for_differential: "; c=csmc.cliqId, is_root=is_root, any_init=any_init)
          linear_on_manifold && !is_root && !any_init
        end

        if init_for_differential
          frontal_vars = getVariable.(csmc.cliqSubFg,  getCliqFrontalVarIds(csmc.cliq))
          filter!(!isInitialized, frontal_vars)
          foreach(fvar->getSolverData(fvar, csmc.solveKey).initialized = true, frontal_vars)
          logCSM(csmc, "CSM-2b init_for_differential: "; c=csmc.cliqId,lbl=getLabel.(frontal_vars))
        end
        ## END experimental

    setCliqueDrawColor!(csmc.cliq, "green")

    logCSM(csmc, "CSM-2b Trying up init -- all not initialized"; c=csmc.cliqId)

    # structure for all up message densities computed during this initialization procedure.
    varorder = getCliqVarInitOrderUp(csmc.cliqSubFg)
    someInit = cycleInitByVarOrder!(csmc.cliqSubFg, varorder, solveKey=csmc.solveKey, logger=csmc.logger)
    # is clique fully upsolved or only partially?
    # print out the partial init status of all vars in clique
    printCliqInitPartialInfo(csmc.cliqSubFg, csmc.cliq, csmc.solveKey, csmc.logger)
    logCSM(csmc, "CSM-2b solveUp try init -- someInit=$someInit, varorder=$varorder"; c=csmc.cliqId)

    someInit ? setCliqueDrawColor!(csmc.cliq, "darkgreen") :  setCliqueDrawColor!(csmc.cliq, "lightgreen")

    solveStatus = someInit ? INITIALIZED : NO_INIT

        ## FIXME init to whatever is in frontals
        # set frontals init back to false
        if init_for_differential #experimental_sommer_init_to_whatever_is_in_frontals
          foreach(fvar->getSolverData(fvar, csmc.solveKey).initialized = false, frontal_vars)
          if someInit
            solveStatus = UPSOLVED
          end
        end
        ## END EXPERIMENTAL

  setCliqueStatus!(csmc.cliq, solveStatus)

  return postUpSolve_StateMachine

end

"""
  $SIGNATURES

Notes
- State machine function 2c
"""
function solveUp_StateMachine(csmc::CliqStateMachineContainer)

  logCSM(csmc, "CSM-2c, cliq=$(csmc.cliqId) Solving Up")

  setCliqueDrawColor!(csmc.cliq, "red")

  #Make sure all are initialized
  if !areCliqVariablesAllInitialized(csmc.cliqSubFg, csmc.cliq, csmc.solveKey)
    logCSM(csmc, "CSM-2c All children upsolved, not init, try init then upsolve"; c=csmc.cliqId)
    varorder = getCliqVarInitOrderUp(csmc.cliqSubFg)
    someInit = cycleInitByVarOrder!(csmc.cliqSubFg, varorder, solveKey=csmc.solveKey, logger=csmc.logger)
  end

  isinit = areCliqVariablesAllInitialized(csmc.cliqSubFg, csmc.cliq, csmc.solveKey)
  logCSM(csmc, "CSM-2c midway, isinit=$isinit")
  # Check again
  if isinit
    logCSM(csmc, "CSM-2c doing upSolve -- all initialized")

    __doCliqUpSolveInitialized!(csmc)

    setCliqueStatus!(csmc.cliq, UPSOLVED)

  else
    _dbgCSMSaveSubFG(csmc, "fg_child_solved_cant_init")
    # it can be a leaf
    logCSM(csmc, "CSM-2c solveUp -- all children upsolved, but init failed.")
  end

  # if converged_and_happy

  # else # something went wrong propagate error
  #   @error "X-3, something wrong with solve up"
  #   # propagate error to cleanly exit all cliques
  #   putErrorUp(csmc)
  #   if length(getParent(csmc.tree, csmc.cliq)) == 0
  #     putErrorDown(csmc)
  #     return IncrementalInference.exitStateMachine
  #   end

  #   return waitForDown_StateMachine
  # end

  return postUpSolve_StateMachine
end

"""
  $SIGNATURES

CSM function only called when `getSolverParams(dfg).upsolve == false` that tries to skip upsolve.
Notes
- Cliques are uprecycled to add differential messages.
- State machine function 2d
"""
function tryDownSolveOnly_StateMachine(csmc::CliqStateMachineContainer)
  logCSM(csmc, "CSM-2d tryDownSolveOnly_StateMachine clique $(csmc.cliqId) status $(getCliqueStatus(csmc.cliq))")

  logCSM(csmc, "CSM-2d Skipping upsolve clique $(csmc.cliqId)"; loglevel=Logging.Info, st=getCliqueStatus(csmc.cliq))
  if getCliqueStatus(csmc.cliq) == NULL
    logCSM(csmc, "CSM-2d Clique $(csmc.cliqId) status NULL, trying as UPRECYCLED"; loglevel=Logging.Warn)

    # Are all variables solved at least once?
    if all(getSolvedCount.(getVariables(csmc.cliqSubFg)) .> 0)
      setCliqueStatus!(csmc.cliq, UPRECYCLED)
    else
      logCSM(csmc, "CSM-2d Clique $(csmc.cliqId) cannot be UPRECYCLED, all variables not solved. Set solverParams to upsolve=true.";
              loglevel=Logging.Error)
      # propagate error to cleanly exit all cliques
      putErrorUp(csmc)
      if length(getParent(csmc.tree, csmc.cliq)) == 0
        putErrorDown(csmc)
        return IncrementalInference.exitStateMachine
      end
      return waitForDown_StateMachine

    end

  end

  return postUpSolve_StateMachine
end

"""
    $SIGNATURES

Post-upsolve remove message factors and send messages
Notes
- State machine function 2e
"""
function postUpSolve_StateMachine(csmc::CliqStateMachineContainer)

  solveStatus = getCliqueStatus(csmc.cliq)
  # fill in belief
  logCSM(csmc, "CSM-2e prepCliqueMsgUp, going for prepCliqueMsgUp")
  beliefMsg = prepCliqueMsgUp(csmc.cliqSubFg, csmc.cliq, csmc.solveKey,
                              solveStatus, logger=csmc.logger,
                              sender=(; id=csmc.cliq.id.value,
                                        step=csmc._csm_iter) )
  #

  logCSM(csmc, "CSM-2e prepCliqueMsgUp", msgon=keys(beliefMsg.belief), beliefMsg=beliefMsg)

  # Done with solve delete factors
  # remove msg factors that were added to the subfg
  tags_ = getSolverParams(csmc.cliqSubFg).useMsgLikelihoods ? [:__UPWARD_COMMON__;] : [:__LIKELIHOODMESSAGE__;]
  msgfcts= deleteMsgFactors!(csmc.cliqSubFg, tags_)
  logCSM(csmc, "CSM-2e doCliqUpsSolveInit.! -- status = $(solveStatus), removing $(tags_) factors, length=$(length(msgfcts))")

  # store the cliqSubFg for later debugging
  _dbgCSMSaveSubFG(csmc, "fg_afterupsolve")

  # warn and clean exit on stalled tree init
  if csmc.init_iter > getSolverParams(csmc.cliqSubFg).limittreeinit_iters
    logCSM(csmc, "CSM-2e Clique $(csmc.cliqId) tree init failed, max init retries reached."; loglevel=Logging.Error)
    putErrorUp(csmc)
    if length(getParent(csmc.tree, csmc.cliq)) == 0
      putErrorDown(csmc)
      return IncrementalInference.exitStateMachine
    end
    return waitForDown_StateMachine
  end

  # always put up belief message in upTx, only used for debugging isolated cliques
  getMessageBuffer(csmc.cliq).upTx = deepcopy(beliefMsg)
  #propagate belief
  for e in getEdgesParent(csmc.tree, csmc.cliq)
    logCSM(csmc, "CSM-2e $(csmc.cliq.id): put! on edge $(e)")
    putBeliefMessageUp!(csmc.tree, e, beliefMsg)
  end

  if getSolverParams(csmc.dfg).downsolve
    return waitForDown_StateMachine
  else
    return updateFromSubgraph_StateMachine
  end
end

## =========================================================================================
## Wait for Down -- 3
## =========================================================================================

"""
    $SIGNATURES

Notes
- State machine function waitForDown 3
"""
function waitForDown_StateMachine(csmc::CliqStateMachineContainer)

  logCSM(csmc, "CSM-3 wait for down messages if needed")

  # setCliqueDrawColor!(csmc.cliq, "lime")

  for e in getEdgesParent(csmc.tree, csmc.cliq)
    logCSM(csmc, "CSM-3 $(csmc.cliq.id): take! on edge $(e)")
    # Blocks until data is available.
    beliefMsg = takeBeliefMessageDown!(csmc.tree, e) # take!(csmc.tree.messageChannels[e.index].downMsg)
    logCSM(csmc, "CSM-3 $(csmc.cliq.id): Belief message received with status $(beliefMsg.status)")

    logCSM(csmc, "CSM-3 down msg on $(keys(beliefMsg.belief))"; beliefMsg=beliefMsg)
    # save down incoming message for use and debugging
    getMessageBuffer(csmc.cliq).downRx = beliefMsg

    # Down branching happens here

    # ERROR_STATUS
    if beliefMsg.status == ERROR_STATUS
      putErrorDown(csmc)
      return IncrementalInference.exitStateMachine

    elseif csmc.algorithm == :parametric
      beliefMsg.status != DOWNSOLVED && error("#FIXME")
      return solveDown_ParametricStateMachine
    elseif beliefMsg.status in [MARGINALIZED, DOWNSOLVED, INITIALIZED, NO_INIT]
      return preDownSolve_StateMachine
    # elseif beliefMsg.status == DOWNSOLVED
    #   return solveDown_StateMachine
    # elseif beliefMsg.status == INITIALIZED || beliefMsg.status == NO_INIT
    #   return tryDownInit_StateMachine
    else
      logCSM(csmc, "CSM-3 Unknown state"; status=beliefMsg.status, loglevel=Logging.Error, c=csmc.cliqId)
      error("CSM-3 waitForDown State Error: Unknown/unimplemented transision.")
    end
  end

  # The clique is a root
  # root clique down branching happens here
  if csmc.algorithm == :parametric
    return solveDown_ParametricStateMachine
  else
    return preDownSolve_StateMachine
  end

end

## =========================================================================================
## Down Functions -- 4
## =========================================================================================

## TODO Consolidate
function CliqDownMessage(csmc::CliqStateMachineContainer, status=DOWNSOLVED)

  #JT TODO maybe use Tx buffer
  newDwnMsgs = LikelihoodMessage( sender=(; id=csmc.cliq.id.value,
                                            step=csmc._csm_iter),
                                  status=status)

  # create all messages from subfg
  for mk in getCliqFrontalVarIds(csmc.cliq)
    v = getVariable(csmc.cliqSubFg, mk)
    if isInitialized(v, csmc.solveKey)
      newDwnMsgs.belief[mk] = TreeBelief(v, csmc.solveKey)
    end
  end

  logCSM(csmc, "cliq $(csmc.cliq.id), CliqDownMessage, allkeys=$(keys(newDwnMsgs.belief))")

  return newDwnMsgs
end

"""
    $SIGNATURES

Notes
- State machine function 4a
"""
function preDownSolve_StateMachine(csmc::CliqStateMachineContainer)
  logCSM(csmc, "CSM-4a Preparing for down init/solve")

  opts = getSolverParams(csmc.dfg)

  # get down msg from Rx buffer (saved in take!)
  dwnmsgs = getMessageBuffer(csmc.cliq).downRx

  # DownSolve cliqSubFg
  #only down solve if its not a root and not MARGINALIZED
  if length(getParent(csmc.tree, csmc.cliq)) != 0 && getCliqueStatus(csmc.cliq) != MARGINALIZED

    logCSM(csmc, "CSM-4a doCliqDownSolve_StateMachine -- dwnmsgs=$(collect(keys(dwnmsgs.belief)))")
    # maybe cycle through separators (or better yet, just use values directly -- see next line)
    msgfcts = addMsgFactors!(csmc.cliqSubFg, dwnmsgs, DownwardPass)
    # force separator variables in cliqSubFg to adopt down message values
    updateSubFgFromDownMsgs!(csmc.cliqSubFg, dwnmsgs, getCliqSeparatorVarIds(csmc.cliq))

    if dwnmsgs.status in [DOWNSOLVED, MARGINALIZED]
      logCSM(csmc, "CSM-4a doCliqDownSolve_StateMachine")
      return solveDown_StateMachine
    elseif dwnmsgs.status == INITIALIZED || dwnmsgs.status == NO_INIT
      return tryDownInit_StateMachine
    else
      logCSM(csmc, "CSM-4a Unknown state"; status=dwnmsgs.status, loglevel=Logging.Error, c=csmc.cliqId)
      error("CSM-4a waitForDown State Error: Unknown/unimplemented transision.")
    end
  else
    # Special root case or MARGINALIZED
    #TODO improve
    solveStatus = getCliqueStatus(csmc.cliq)
    logCSM(csmc, "CSM-4a root case or MARGINALIZED"; status=solveStatus, c=csmc.cliqId)
    if solveStatus in [INITIALIZED, NO_INIT, UPSOLVED, UPRECYCLED, MARGINALIZED]
      solveStatus == MARGINALIZED &&  setCliqueDrawColor!(csmc.cliq, "blue")
      if solveStatus in [UPSOLVED, UPRECYCLED]
        setCliqueStatus!(csmc.cliq, DOWNSOLVED)
      end
      return postDownSolve_StateMachine
    else
      error("CSM-4a unknown status root $solveStatus")
    end
  end

end

"""
    $SIGNATURES

Notes
- State machine function 4b
"""
function tryDownInit_StateMachine(csmc::CliqStateMachineContainer)

  setCliqueDrawColor!(csmc.cliq, "olive")

  logCSM(csmc, "CSM-4b Trying Down init -- all not initialized")

  # structure for all up message densities computed during this initialization procedure.
  # XXX
  dwnkeys_ = lsf(csmc.cliqSubFg, tags=[:__DOWNWARD_COMMON__;]) .|> x->ls(csmc.cliqSubFg, x)[1]
  initorder = getCliqInitVarOrderDown(csmc.cliqSubFg, csmc.cliq, dwnkeys_)
  # initorder = getCliqVarInitOrderUp(csmc.tree, csmc.cliq)

  someInit = cycleInitByVarOrder!(csmc.cliqSubFg, initorder, solveKey=csmc.solveKey, logger=csmc.logger)
  # is clique fully upsolved or only partially?
  # print out the partial init status of all vars in clique
  printCliqInitPartialInfo(csmc.cliqSubFg, csmc.cliq, csmc.solveKey, csmc.logger)
  logCSM(csmc, "CSM-4b tryInitCliq_StateMachine -- someInit=$someInit, varorder=$initorder")

  msgfcts = deleteMsgFactors!(csmc.cliqSubFg, [:__DOWNWARD_COMMON__;]) # msgfcts # TODO, use tags=[:__LIKELIHOODMESSAGE__], see #760
  logCSM(csmc, "CSM-4b tryDownInit_StateMachine - removing factors, length=$(length(msgfcts))")

  solveStatus = someInit ? INITIALIZED : NO_INIT
  someInit ? setCliqueDrawColor!(csmc.cliq, "seagreen") :  setCliqueDrawColor!(csmc.cliq, "khaki")
  setCliqueStatus!(csmc.cliq, solveStatus)

  return postDownSolve_StateMachine
end

"""
    $SIGNATURES

Notes
- State machine function 4c
"""
function solveDown_StateMachine(csmc::CliqStateMachineContainer)

  logCSM(csmc, "CSM-4c Solving down")

  setCliqueDrawColor!(csmc.cliq, "maroon")

  # DownSolve cliqSubFg
  # add messages, do downsolve, remove messages

  # get down msg from Rx buffer (saved in take!)
  dwnmsgs = getMessageBuffer(csmc.cliq).downRx
  # #XXX
  # # maybe cycle through separators (or better yet, just use values directly -- see next line)
  # msgfcts = addMsgFactors!(csmc.cliqSubFg, dwnmsgs, DownwardPass)
  # force separator variables in cliqSubFg to adopt down message values
  # updateSubFgFromDownMsgs!(csmc.cliqSubFg, dwnmsgs, getCliqSeparatorVarIds(csmc.cliq))

  opts = getSolverParams(csmc.dfg)
  #XXX test with and without
  # add required all frontal connected factors
  if !opts.useMsgLikelihoods
    newvars, newfcts = addDownVariableFactors!(csmc.dfg, csmc.cliqSubFg, csmc.cliq, csmc.logger, solvable=1)
  end

  # store the cliqSubFg for later debugging
  _dbgCSMSaveSubFG(csmc, "fg_beforedownsolve")

  ## new way
  # calculate belief on each of the frontal variables and iterate if required
  solveCliqDownFrontalProducts!(csmc.cliqSubFg, csmc.cliq, opts, csmc.logger)
  csmc.dodownsolve = false

  logCSM(csmc, "CSM-4c solveDown -- finished with downGibbsCliqueDensity, now update csmc")

  # update clique subgraph with new status
  # setCliqueDrawColor!(csmc.cliq, "lightblue")

  # remove msg factors that were added to the subfg
  rmFcts = deleteMsgFactors!(csmc.cliqSubFg)
  logCSM(csmc, "CSM-4c solveDown -- removing up message factors, length=$(length(rmFcts))")

  # store the cliqSubFg for later debugging
  _dbgCSMSaveSubFG(csmc, "fg_afterdownsolve")

  setCliqueStatus!(csmc.cliq, DOWNSOLVED)

  logCSM(csmc, "CSM-4c $(csmc.cliq.id): clique down solve completed")

  return postDownSolve_StateMachine

end

"""
    $SIGNATURES

Notes
- State machine function 4d
"""
function postDownSolve_StateMachine(csmc::CliqStateMachineContainer)

  solveStatus = getCliqueStatus(csmc.cliq)
  #fill in belief
  #TODO use prepSetCliqueMsgDownConsolidated
  beliefMsg = CliqDownMessage(csmc, solveStatus)

  if length(keys(beliefMsg.belief)) == 0
    logCSM(csmc, "CSM-4d Empty message on clique $(csmc.cliqId) frontals"; loglevel=Logging.Info)
  end

  logCSM(csmc, "CSM-4d msg to send down on $(keys(beliefMsg.belief))"; beliefMsg=beliefMsg)
  # pass through the frontal variables that were sent from above
  downmsg = getMessageBuffer(csmc.cliq).downRx
  svars = getCliqSeparatorVarIds(csmc.cliq)
  if !isnothing(downmsg)
    pass_through_separators = intersect(svars, keys(downmsg.belief))
    for si in pass_through_separators
      beliefMsg.belief[si] = downmsg.belief[si]
      logCSM(csmc, "CSM-4d adding parent message"; sym=si, msg=downmsg.belief[si])
    end
  end

  # Store the down message for debugging, will be stored even if no children present
  getMessageBuffer(csmc.cliq).downTx = beliefMsg

  #TODO maybe send a specific message to only the child that needs it
  @sync for e in getEdgesChildren(csmc.tree, csmc.cliq)
    logCSM(csmc, "CSM-4d $(csmc.cliq.id): put! on edge $(e)")
    @async putBeliefMessageDown!(csmc.tree, e, beliefMsg)#put!(csmc.tree.messageChannels[e.index].downMsg, beliefMsg)
  end

  if getCliqueStatus(csmc.cliq) in [DOWNSOLVED, MARGINALIZED]

    return updateFromSubgraph_StateMachine

  else
    # detete all message factors to start clean
    deleteMsgFactors!(csmc.cliqSubFg)

    return waitForUp_StateMachine
  end

end

## =========================================================================================
## Finalize Functions -- 5
## =========================================================================================

"""
    $SIGNATURES

The last step in CSM to update the main FG from the sub FG.

Notes
- CSM function 5
"""
function updateFromSubgraph_StateMachine(csmc::CliqStateMachineContainer)

  isParametricSolve = csmc.algorithm == :parametric

  # set PPE and solved for all frontals
  if !isParametricSolve
    for sym in getCliqFrontalVarIds(csmc.cliq)
      # set PPE in cliqSubFg
      setVariablePosteriorEstimates!(csmc.cliqSubFg, sym)
      # set solved flag
      vari = getVariable(csmc.cliqSubFg, sym, csmc.solveKey)
      setSolvedCount!(vari, getSolvedCount(vari, csmc.solveKey)+1, csmc.solveKey )
    end
  end

  # transfer results to main factor graph
  frsyms = getCliqFrontalVarIds(csmc.cliq)
  logCSM(csmc, "CSM-5 finishingCliq -- transferUpdateSubGraph! with solveKey=$(csmc.solveKey) on $frsyms")
  transferUpdateSubGraph!(csmc.dfg, csmc.cliqSubFg, frsyms, csmc.logger, solveKey=csmc.solveKey, updatePPE=!isParametricSolve)

  #solve finished change color
  setCliqueDrawColor!(csmc.cliq, "lightblue")

  logCSM(csmc, "CSM-5 Clique $(csmc.cliq.id) finished, solveKey=$(csmc.solveKey)", loglevel=Logging.Info)
  return IncrementalInference.exitStateMachine
end

#

ext: jl


text:
if !isdefined(:KernelFunctions); include("KernelFunctions.jl"); end;
if !isdefined(:CustomKMeans); include("AFKMC2.jl"); end;

module BayesianSVM

using KernelFunctions
using CustomKMeans
using Distributions
using StatsBase
using PyPlot

export BSVM

#Corresponds to the BSVM model
type BSVM
  #Data
  X #Feature vectors
  y #Labels (-1,1)
  #Stochastic parameters
  Stochastic::Bool
    nSamplesUsed::Int64 #Size of the minibatch used
    κ_s::Float64 #Parameters for decay of learning rate (iter + κ)^-τ in case adaptative learning rate is not used
    τ_s::Float64
  #Non linear parameters
  NonLinear::Bool
  Sparse::Bool
    kernels::Array{Kernel,1} #Kernels function used
    γ::Float64 # Regularization parameter of the noise
    m::Int64 #Number of inducing points
  #Autotuning parameters
  Autotuning::Bool
    κ_Θ::Float64 #Parameters for decay of learning rate for the hyperparameter  (iter + κ)^-τ
    τ_Θ::Int64
    autotuningFrequency::Int64 #Frequency of update of the hyperparameter
  #Flag for adaptative learning rate for the SVI
  AdaptativeLearningRate::Bool
  #General Parameters for training
  Intercept::Bool
  ϵ::Float64 #Desired precision (on ||β(t)-β(t-1)||)
  nEpochs::Int64 #Maximum number of iterations
  β_init::Array{Float64,1} #Initial value for β
  smoothingWindow::Int64
  VerboseLevel::Int64
  Storing::Bool #Store values for debugging and visualization purposes
    StoringFrequency::Int64 #Every X steps
    StoredValues::Array{Float64,2}
    StoredDerivativesELBO::Array{Float64,2}
  #Functions
  Kernel_function::Function #kernel function associated with the model
  Train::Function #Model train for a certain number of iterations
  Predict::Function
  PredictProba::Function
  ELBO::Function
  DerivativesELBO::Function
  Plotting::Function
  Update::Function

  #Parameters learned with training
  nSamples::Int64 # Number of data points
  nFeatures::Int64 # Number of features
  μ::Array{Float64,1} # Mean for variational distribution
  η_1::Array{Float64,1} #Natural Parameter #1
  ζ::Array{Float64,2} # Covariance matrix of variational distribution
  η_2::Array{Float64,2} #Natural Parameter #2
  α::Array{Float64,1} # Distribution parameter of the GIG distribution of the latent variables
  invΣ::Array{Float64,2} #Inverse Prior Matrix for the linear case
  invK::Array{Float64,2} #Inverse Kernel Matrix for the nonlinear case
  invKmm::Array{Float64,2} #Inverse Kernel matrix of inducing points
  Ktilde::Array{Float64,1} #Diagonal of the covariance matrix between inducing points and generative points
  κ::Array{Float64,2} #Kmn*invKmm
  inducingPoints::Array{Float64,2} #Inducing points coordinates for the Big Data GP
  top #Storing matrices for repeated predictions (top and down are numerator and discriminator)
  down
  MatricesPrecomputed::Bool #Flag to know if matrices needed for predictions are already computed or not
  ρ_s::Float64 #Learning rate for CAVI
  g::Array{Float64,1} # g & h are expected gradient value for computing the adaptive learning rate
  h::Float64
  ρ_Θ::Float64 # learning rate for auto tuning
  initialized::Bool
  evol_β::Array{Float64,2} #Store the betas for smooth convergence criterium

  #Constructor
  function BSVM(X,y;Stochastic::Bool=false,Sparse::Bool=true,NonLinear::Bool=true,AdaptativeLearningRate::Bool=true,Autotuning::Bool=false,
                                  nEpochs::Int64 = 2000,batchSize::Int64=-1,κ_s::Float64=1.0,τ_s::Int64=100,
                                  kernels=0,γ::Float64=1e-3,m::Int64=0,κ_Θ::Float64=1.0,τ_Θ::Int64=100,autotuningfrequency::Int64=10,
                                  Intercept::Bool=false,ϵ::Float64=1e-5,β_init=[0.0],smoothingWindow::Int64=10,
                                  Storing::Bool=false,StoringFrequency::Int64=1,VerboseLevel::Int64=0)
    iter = 1
    if kernels == 0 && NonLinear
      warn("No kernel indicated, a rbf kernel function with lengthscale 1 is used")
      kernels = [Kernel("rbf",1.0,params=1.0)]
    end
    #Verification of consistency of the model
    this = new(X,y,Stochastic,batchSize,κ_s,τ_s,NonLinear,Sparse,kernels,γ,m,Autotuning,κ_Θ,τ_Θ,autotuningfrequency,AdaptativeLearningRate,Intercept,ϵ,nEpochs,β_init,smoothingWindow,VerboseLevel,Storing,StoringFrequency)
    if !ModelVerification(this,size(this.X),size(this.y))
      return
    end
    this.initialized = false
    if NonLinear
      this.top = 0
      this.down = 0
      MatricesPrecomputed = false
      this.Kernel_function = function(X1,X2)
          dist = 0
          for i in 1:size(this.kernels,1)
            dist += this.kernels[i].coeff*this.kernels[i].compute(X1,X2)
          end
          return dist
      end
      this.Train = function()
          TrainBSVM(this);
        end

      if Sparse
        this.Predict = function(X_test)
            SparsePredict(X_test,this)
          end
        this.PredictProba = function(X_test)
            SparsePredictProb(X_test,this)
          end
         #if this.Stochastic
        #     this.ELBO = function(X,y)
        #         NoisySparseELBO(this,y)
        #       end
        #  else
            this.ELBO = function()
                SparseELBO(this)
              end
        #   end
          this.DerivativesELBO = function() #Not implemented yet
              return 0
          end
      else
        this.Predict = function(X_test)
            NonLinearPredict(X_test,this)
          end
        this.PredictProba = function(X_test)
            NonLinearPredictProb(X_test,this)
          end
        this.ELBO = function()
            NonLinearELBO(this)
          end
        this.DerivativesELBO = function()
            DerivativesNonLinearELBO(this.y,this.μ,this.ζ,this.α,this.invK,this.Autotuning ? this.J : eye(size(this.X,1)))
          end
      end

    else
      this.Predict = function(X_test)
          LinearPredict(X_test,this)
        end
      this.PredictProba = function(X_test)
          LinearPredictProb(X_test,this)
        end
      this.ELBO = function()
          LinearELBO(this)
        end
      this.DerivativesELBO = function()
          DerivativesLinearELBO(Diagonal(this.y)*this.X,this.μ,this.ζ,this.α,inv(this.invΣ))
        end
    end
    this.Plotting = function(;s::String="All")
        ParametersPlotting(this,option=s)
      end
    this.Update = function(iter)
        Update(this,this.X,this.y,iter)
      end
    return this
  end
  #end of constructor
end

#Function to check consistency of the different parameters and the possible correction of some of them in some cases
function ModelVerification(model::BSVM,XSize,ySize)
  if model.Intercept && model.NonLinear
    warn("Not possible to have intercept for the non linear case, removing automatically this option")
    model.Intercept = false
  end
  if XSize[1] != ySize[1]
    warn("There is a dimension problem with the data size(y) != size(X)")
    return false
  end
  if model.NonLinear && model.Sparse
      minpoints = 50
    if model.m > XSize[1]
        if XSize[1] < 3.0*minpoints
            warn("Wrong inducing point setting, using the full batch method")
            model.Sparse = false
            model.Stochastic = false
        else
            warn("There are more inducing points than actual points, setting it to 10% of the datasize (minimum of $minpoints points)")
            model.m = min(minpoints,XSize[1]÷10)
        end
    elseif model.m == 0
        if XSize[1] < 3.0*minpoints
            warn("Number of inducing points was not manually set, using the full batch method")
            model.Sparse = false
            model.Stochastic = false
        else
            warn("Number of inducing points was not manually set, setting it to 10% of the datasize (minimum of $minpoints points)")
            model.m = min(minpoints,XSize[1]÷10)
        end
    end
  end
  if !model.Stochastic && model.smoothingWindow > 1
      # warn("Smoothing Windows larger than one is unnecessary for a non-stochastic optimization, setting it to one.")
      model.smoothingWindow = 1;
  end
  if model.Sparse && !model.NonLinear
    warn("Model cannot be sparse and linear at the same time, assuming linear model")
    model.Sparse = false;
  end
  if model.γ <= 0
    warn("Gamma should be strictly positive, setting it to default value 1.0e-3")
    model.γ = 1e-3
  end
  if model.nSamplesUsed == -1 && model.Stochastic
    warn("No batch size has been given, stochastic option has been removed")
    model.Stochastic = false
  end
  return true
end

function TrainBSVM(model::BSVM)
  model.nSamples = length(model.y)
  model.nFeatures = model.NonLinear ? (model.Sparse ? model.m : length(model.y)) : size(model.X,2)

  if model.VerboseLevel > 0
    println("Starting training of data of size $((model.nSamples,size(model.X,2))), using the"*(model.Autotuning ? " autotuned" : "")*(model.Stochastic ? " stochastic" : "")*(model.NonLinear ? " kernel" : " linear")*" method"
    *(model.AdaptativeLearningRate ? " with adaptative learning rate" : "")*(model.Sparse ? " with inducing points" : ""))
  end

  #Initialization of the variables
  if model.Intercept
    model.nFeatures += 1
    model.X = [ones(Float64,model.nSamples) model.X]
  end
  if !model.initialized
    if model.β_init[1] == 0 || length(model.β_init) != nFeatures
      if model.VerboseLevel > 2
        warn("Initial mean of the variational distribution is sampled from a multinormal distribution")
      end
      model.μ = randn(model.nFeatures)
    else
      model.μ = model.β_init
    end
    #Necessary to initialize only for first computation of the ELBO
    model.α = abs.(rand(model.nSamples))
    model.ζ = eye(model.nFeatures)
    #Creation of the Kernel Matrix and its inverse in the different cases as well as the prior
    if model.NonLinear
      if !model.Sparse
        model.invK = inv(Symmetric(CreateKernelMatrix(model.X,model.Kernel_function) + model.γ*eye(model.nFeatures),:U))
      end
      if model.Sparse
        model.inducingPoints = KMeansInducingPoints(model.X,model.m,10)
        model.invKmm = Matrix(Symmetric(inv(CreateKernelMatrix(model.inducingPoints,model.Kernel_function)+model.γ*eye(model.nFeatures))))
        Knm = CreateKernelMatrix(model.X,model.Kernel_function,X2=model.inducingPoints)
        model.κ = Knm*model.invKmm
        model.Ktilde = CreateDiagonalKernelMatrix(model.X,model.Kernel_function) + model.γ*ones(model.nSamples) - squeeze(sum(model.κ.*Knm,2),2) #diag(model.κ*transpose(Knm))
      end
    elseif !model.NonLinear
      model.invΣ =  (1.0/model.γ)*eye(model.nFeatures)
    end
    if (model.nSamplesUsed <= 0 || model.nSamplesUsed > model.nSamples)
      model.nSamplesUsed = model.nSamples
    end
    #Initialization of the natural parameters
    model.η_2 = -0.5*inv(model.ζ)
    model.η_1 = 2*model.η_2*model.μ
    if model.AdaptativeLearningRate && model.Stochastic
      batchindices = StatsBase.sample(1:model.nSamples,model.nSamplesUsed,replace=false)
      Z = model.NonLinear ? Diagonal(model.y) : Diagonal(model.y)*model.X
      (grad_1,grad_2) = NaturalGradientELBO(model.α,(model.NonLinear && model.Sparse) ? Z*model.κ : Z, model.NonLinear ? (model.Sparse ? model.invKmm : model.invK) : model.invΣ,model.nSamples/model.nSamplesUsed)
      model.τ_s = model.nSamplesUsed
      model.g = vcat(grad_1,reshape(grad_2,size(grad_2,1)^2))
      model.h = norm(vcat(grad_1,reshape(grad_2,size(grad_2,1)^2)))^2
    end
    model.ρ_Θ = model.Autotuning? (1+model.τ_Θ)^(-model.κ_Θ) : 1.0;
    model.ρ_s = model.Stochastic ? (model.AdaptativeLearningRate ? dot(model.g,model.g)/model.h : (1+model.τ_s)^(-model.κ_s)) : 1.0
    if model.Storing
      # Storing trace(ζ),ELBO,max(|α|),ρ_s,ρ_Θ/ρ_γ,||Θ||/γ
      model.StoredValues = zeros(model.nEpochs,6)
      model.StoredDerivativesELBO = zeros(model.nEpochs,4)
      model.StoredValues[1,:] = [model.ELBO(),trace(model.ζ),mean(model.μ),mean(model.α),model.γ,model.Autotuning ? model.ρ_Θ : 0.0]
      model.StoredDerivativesELBO[1,:] = model.DerivativesELBO()
    end
    model.initialized = true
    model.down = 0
    model.top = 0
    model.MatricesPrecomputed = false
  end
  evol_β = zeros(model.nEpochs,model.nFeatures)
  evol_β[1,:] = model.μ
  evol_ELBO = zeros(model.nEpochs)
  evol_ELBO[1] = model.ELBO()

  batchindices = collect(1:model.nSamples)
  current = evol_ELBO[1]
  conv = Inf #Initialization of the Convergence value
  iter::Int64 = 1
  ##End of Initialization of the parameters
  if model.VerboseLevel > 1
    println("Iteration $iter / $(model.nEpochs) (max)")
    println("Convergence : $conv, ELBO : $current")
  end
  #Two criterions for stopping, number of iterations or convergence
  while iter < model.nEpochs && conv > model.ϵ
    #Print some of the parameters
    model.Update(iter)
    iter += 1
    evol_β[iter,:] = model.μ
    evol_ELBO[iter] = model.ELBO()
    smooth_1 = mean(evol_ELBO[max(1,iter-2*model.smoothingWindow):iter-1,:]);smooth_2 = mean(evol_ELBO[max(2,iter-2*model.smoothingWindow+1):iter,:]);
    conv = abs.(smooth_1-smooth_2)/abs.(smooth_1)
    #smooth_1 = mean(evol_β[max(1,iter-2*model.smoothingWindow):iter-1,:],1);smooth_2 = mean(evol_β[max(2,iter-2*model.smoothingWindow+1):iter,:],1);
    #conv = norm(smooth_1/norm(smooth_1)-smooth_2/norm(smooth_2))

    #prev = current
    # if model.VerboseLevel > 2 || (model.VerboseLevel > 1  && iter%100==0)
        current_ELBO = evol_ELBO[iter]
    # end
    #conv = abs(current-prev)
    if model.Storing && iter%model.StoringFrequency == 0
      if model.NonLinear && model.Stochastic && ((!model.Autotuning && iter<=2) || (model.Autotuning && ((iter-1)%model.autotuningFrequency == 0)))
        println("Recomputing Kernel matrices")
        model.invK = Matrix(Symmetric(inv(CreateKernelMatrix(model.X,model.Kernel_function)+model.γ*eye(model.nSamples)),:U))
        if model.Autotuning
          model.J = CreateKernelMatrix(model.X,deriv_rbf,model.Θ)
        end
      end
      model.StoredValues[iter÷model.StoringFrequency,:] = [model.ELBO(),trace(model.ζ),mean(model.μ),mean(model.α),model.γ,model.Autotuning ? model.ρ_Θ : 0.0,]
      model.StoredDerivativesELBO[iter÷model.StoringFrequency,:] = model.DerivativesELBO()
      #println(model.StoredDerivativesELBO[iter÷model.StoringFrequency,:])
    end
    if model.VerboseLevel > 2 || (model.VerboseLevel > 1  && iter%100==0)
      println("Iteration $iter / $(model.nEpochs) (max)")
      println("Convergence : $conv, ELBO : $current_ELBO")
      if model.Autotuning
        println("Gamma : $(model.γ)")
        for i in 1:size(model.kernels,1)
          println("(Coeff,Parameter) for kernel $i : $((model.kernels[i].coeff,(model.kernels[i].Nparams > 0)? model.kernels[i].param : 0))")
        end
        println("rho theta : $(model.ρ_Θ)")
      end
    end
  end
  if model.VerboseLevel > 0
    println("Training ended after $iter iterations")
  end
  if model.Storing
    model.StoredValues = model.StoredValues[1:iter÷model.StoringFrequency,:];
    model.StoredDerivativesELBO = model.StoredDerivativesELBO[1:iter÷model.StoringFrequency,:];
    model.evol_β = evol_β[1:iter,:]
  end
  return model
end

function Update(model::BSVM,X::Array{Float64,2},y::Array{Float64,1},iter::Int64) #Coordinates ascent of the parameters
    if model.Stochastic
      batchindices = StatsBase.sample(1:model.nSamples,model.nSamplesUsed,replace=false)
    else
      batchindices = collect(1:model.nSamples)
    end
    model.top = 0; model.down = 0; model.MatricesPrecomputed = false;#Need to recompute the matrices
    #Definition of the Z matrix, different for everycase
    Z = model.NonLinear ? (model.Sparse ? Diagonal(y[batchindices])*model.κ[batchindices,:] : Diagonal(y[batchindices]) ) : (Diagonal(y[batchindices])*X[batchindices,:])
    #Computation of latent variables
    model.α[batchindices] = (1 - Z*model.μ).^2
    model.α[batchindices] +=  squeeze(sum((Z*model.ζ).*Z,2),2)
    if model.Sparse && model.NonLinear
      model.α[batchindices] += model.Ktilde[batchindices] #Cf derivation of updates
    end

    #Compute the natural gradient
    (grad_η_1,grad_η_2) = NaturalGradientELBO(model.α[batchindices],Z, model.NonLinear ? (model.Sparse ? model.invKmm : model.invK) : model.invΣ, model.Stochastic ? model.nSamples/model.nSamplesUsed : 1.0)

    #Compute the learning rate
    if model.AdaptativeLearningRate && model.Stochastic
      #Using the paper on the adaptive learning rate for the SVI (update from the natural gradients)
      model.g = (1-1/model.τ_s)*model.g + vcat(grad_η_1-model.η_1,reshape(grad_η_2-model.η_2,size(grad_η_2,1)^2))./model.τ_s
      model.h = (1-1/model.τ_s)*model.h +norm(vcat(grad_η_1-model.η_1,reshape(grad_η_2-model.η_2,size(grad_η_2,1)^2)))^2/model.τ_s
      model.ρ_s = norm(model.g)^2/model.h
      #if iter%1==0
      #  println("g : $(norm(model.g)^2), h : $(model.h), rho : $(model.ρ_s), tau : $(model.τ_s)")
      #end
      model.τ_s = (1.0 - model.ρ_s)*model.τ_s + 1.0
    elseif model.Stochastic
      #Simple model of learning rate
      model.ρ_s = (iter+model.τ_s)^(-model.κ_s)
    else
      #Non-Stochastic case
      model.ρ_s = 1.0
    end
    model.η_1 = (1.0-model.ρ_s)*model.η_1 + model.ρ_s*grad_η_1; model.η_2 = (1.0-model.ρ_s)*model.η_2 + model.ρ_s*grad_η_2 #Update of the natural parameters with noisy/full natural gradient
    model.ζ = -0.5*inv(model.η_2); model.μ = model.ζ*model.η_1 #Back to the distribution parameters (needed for α updates)

    #Autotuning part, only happens every $autotuningFrequency iterations
    if model.Autotuning && (iter%model.autotuningFrequency == 0)
      if model.VerboseLevel > 2 || (model.VerboseLevel > 1  && iter%100==0)
        println("Before hyperparameter optimization ELBO = $(model.ELBO())")
      end
      model.ρ_Θ = (iter+model.τ_Θ)^(-model.κ_Θ)
      if model.NonLinear
        if model.Sparse
          UpdateHyperparameterSparse!(model)
          model.invKmm = Matrix(Symmetric(inv(CreateKernelMatrix(model.inducingPoints,model.Kernel_function)+model.γ*eye(model.m)),:U))
          model.κ = CreateKernelMatrix(model.X,model.Kernel_function,X2=model.inducingPoints)*model.invKmm
        else
          UpdateHyperparameterNonLinear!(model)
          model.invK = Matrix(Symmetric(inv(CreateKernelMatrix(model.X,model.Kernel_function)+model.γ*eye(model.nSamples)),:U))
        end
      else
        model.γ = UpdateHyperparameterLinear(model)
        model.invΣ = (1/model.γ)*eye(model.nFeatures)
      end
      if model.VerboseLevel > 2 || (model.VerboseLevel > 1  && iter%100==0)
        println("After hyperparameter optimization ELBO = $(model.ELBO())")
      end
    end
end;

function NaturalGradientELBO(α,Z,invPrior,stoch_coef)
  grad_1 =  stoch_coef*transpose(Z)*(1./sqrt.(α)+1)
  grad_2 = -0.5*(stoch_coef*transpose(Z)*Diagonal(1./sqrt.(α))*Z + invPrior)
  (grad_1,grad_2)
end

function UpdateHyperparameterLinear!(model) #Gradient ascent for γ, noise
    model.γ = model.γ + model.ρ_Θ*0.5*((trace(model.ζ)+norm(model.μ))/(model.γ^2.0)-model.nFeatures/model.γ)
end

function UpdateHyperparameterNonLinear!(model)#Gradient ascent for Θ , kernel parameters
    if model.invK == 0
      model.invK = Matrix(Symmetric(inv(CreateKernelMatrix(model.X,model.Kernel_function)+model.γ*eye(model.nSamples)),:U))
    end
    NKernels = size(model.kernels,1)
    A = model.invK*model.ζ-eye(model.nFeatures)
    grad_γ = 0.5*(sum(model.invK.*A)+dot(model.μ,model.invK*model.invK*model.μ))
    if model.VerboseLevel > 2
      println("Grad gamma : $grad_γ")
    end
    model.γ = ((model.γ + model.ρ_Θ*grad_γ) < 0 ) ? model.γ/2 : (model.γ + model.ρ_Θ*grad_γ)
    #Update of both the coefficients and hyperparameters of the kernels
    if NKernels > 1 #If multiple kernels only update the kernel weight
      for i in 1:NKernels
        V = model.invK.*CreateKernelMatrix(model.X,model.kernels[i].compute)
        grad =  0.5*(sum(V.*A)+dot(model.μ,V*model.invK*model.μ))#update of the coeff
        if model.VerboseLevel > 2
          println("Grad kernel $i: $grad")
        end
        model.kernels[i].coeff =  ((model.kernels[i].coeff + model.ρ_Θ*grad) < 0 ) ? model.kernels[i].coeff/2 : (model.kernels[i].coeff + model.ρ_Θ*grad)
      end
    elseif model.kernels[1].Nparams > 0 #If only one update the kernel lengthscale
        V = model.invK*model.kernels[1].coeff*CreateKernelMatrix(model.X,model.kernels[1].compute_deriv)
        grad =  0.5*(sum(V.*A)+dot(model.μ,V*model.invK*model.μ))#update of the hyperparameter
        model.kernels[1].param =  ((model.kernels[1].param + model.ρ_Θ*grad) < 0 ) ? model.kernels[1].param/2 : (model.kernels[1].param + model.ρ_Θ*grad)
        if model.VerboseLevel > 2
          println("Grad kernel: $grad")
        end
    end
end

function UpdateHyperparameterSparse!(model)#Gradient ascent for Θ , kernel parameters #Not finished !!!!!!!!!!!!!!!!!!!!!!!!!!
  NKernels = size(model.kernels,1)
  A = eye(model.nFeatures)-model.invKmm*model.ζ
  B = model.μ*transpose(model.μ) + model.ζ
  Kmn = CreateKernelMatrix(model.inducingPoints,model.Kernel_function;X2=model.X)
  #Computation of noise constant
  Jnm = 0
  ι = (Jnm-model.κ)*model.invKmm
  grad_γ = -0.5*(sum(model.invKmm.*A) - dot(model.μ, transpose(model.μ)*model.invKmm*model.invKmm + 2*transpose(ones(model.nSamples)+1./sqrt.(model.α))*diagm(model.y)*ι)+
  dot(1./sqrt.(model.α),diag(model.κ*(B*transpose(ι)-transpose(Jnm)) + ι*(B*transpose(model.κ)-Kmn) )+ ones(model.nSamples)))
  if model.VerboseLevel > 2
    println("Grad gamma : $grad_γ")
  end
  model.γ = ((model.γ + model.ρ_Θ*grad_γ) < 0 ) ? (model.γ < 1e-7 ? model.γ : model.γ/2) : (model.γ + model.ρ_Θ*grad_γ)
  if NKernels > 1
    for i in 1:NKernels
      Jnm = CreateKernelMatrix(model.X,model.kernels[i].compute,X2=model.inducingPoints)
      Jnn = CreateDiagonalKernelMatrix(model.X,model.kernels[i].compute)
      Jmm = CreateKernelMatrix(model.inducingPoints,model.kernels[i].compute)
      ι = (Jnm-model.κ*Jmm)*model.invKmm
      V = model.invKmm*Jmm
      grad = -0.5*(sum(V.*A) - dot(model.μ, transpose(model.μ)*V*model.invKmm + 2*transpose(ones(model.nSamples)+1./sqrt.(model.α))*diagm(model.y)*ι) +
      dot(1./sqrt.(model.α),diag(model.κ*(B*transpose(ι)-transpose(Jnm)) + ι*(B*transpose(model.κ)-Kmn))+ Jnn))
      model.kernels[i].coeff =  ((model.kernels[i].coeff + model.ρ_Θ*grad) < 0 ) ? model.kernels[i].coeff/2 : (model.kernels[i].coeff + model.ρ_Θ*grad)
      if model.VerboseLevel > 2
        println("Grad kernel $i: $grad")
      end
    end
  elseif model.kernels[1].Nparams > 0 #Update of the hyperparameters of the KernelMatrix
    Jnm = model.kernels[1].coeff*CreateKernelMatrix(model.X,model.kernels[1].compute_deriv,X2=model.inducingPoints)
    Jnn = model.kernels[1].coeff*CreateDiagonalKernelMatrix(model.X,model.kernels[1].compute_deriv)
    Jmm = model.kernels[1].coeff*CreateKernelMatrix(model.inducingPoints,model.kernels[1].compute_deriv)
    ι = (Jnm-model.κ*Jmm)*model.invKmm
    V = model.invKmm*Jmm
    grad = -0.5*(sum(V.*A) - dot(model.μ, transpose(model.μ)*V*model.invKmm + 2*transpose(ones(model.nSamples)+1./sqrt.(model.α))*diagm(model.y)*ι) +
    dot(1./sqrt.(model.α),diag(model.κ*(B*transpose(ι)-transpose(Jnm)) + ι*(B*transpose(model.κ)-Kmn))+Jnn))
    model.kernels[1].param =  ((model.kernels[1].param + model.ρ_Θ*grad) < 0 ) ? model.kernels[1].param/2 : (model.kernels[1].param + model.ρ_Θ*grad)
    if model.VerboseLevel > 2
      println("Grad kernel: $grad, new param is $(model.kernels[1].param)")
    end
  end
end

function LinearELBO(model) #Compute the loglikelihood of the training data, ####-----Could be improved in algebraic form---####
    Z = Diagonal(model.y)*model.X
    ELBO = 0.5*(logdet(model.ζ)+logdet(model.invΣ)-trace(model.invΣ*model.ζ)-dot(model.μ,model.ζ*model.μ));
    for i in 1:model.nSamples
        ELBO += 1.0/2.0*log.(model.α[i]) + log.(besselk.(0.5,model.α[i])) + dot(vec(Z[i,:]),model.μ) + 0.5/model.α[i]*(model.α[i]^2-(1-dot(vec(Z[i,:]),model.μ))^2 - dot(vec(Z[i,:]),model.ζ*vec(Z[i,:])))
    end
    return ELBO
end

function DerivativesLinearELBO(Z,μ,ζ,α,Σ)
    (n,p) = size(Z)
    dζ = 0.5*(inv(ζ)-inv(Σ)-transpose(Z)*Diagonal(1./sqrt.(α))*Z)
    dμ = -inv(ζ)*μ + transpose(Z)*(1./sqrt.(α)+1)
    dα = zeros(n)
    for i in 1:n
      dα[i] = ((1-dot(Z[i,:],μ))^2 + dot(Z[i,:],ζ*Z[i,:]))/(2*(α[i])) - 0.5
    end
    γ = Σ[1,1]
    dγ = 0.5*((trace(ζ)+norm(μ))/(γ^2.0)-p/γ)
    return [trace(dζ),mean(dμ),mean(dα),dγ]
end

function NonLinearELBO(model)
  ELBO = 0.5*(logdet(model.ζ)+logdet(model.invK)-trace(model.invK*model.ζ)-dot(model.μ,model.invK*model.μ))
  for i in 1:model.nSamples
    ELBO += 0.25*log.(model.α[i])+log.(besselk.(0.5,sqrt.(model.α[i])))+model.y[i]*model.μ[i]+(model.α[i]-(1-model.y[i]*model.μ[i])^2-model.ζ[i,i])/(2*sqrt.(model.α[i]))
  end
  return ELBO
end

function DerivativesNonLinearELBO(y,μ,ζ,α,invK,J)
  n = length(y)
  dζ = 0.5*(inv(ζ)-invK-Diagonal(1./sqrt.(α)))
  dμ = -inv(ζ)*μ + Diagonal(y)*(1./sqrt.(α)+1)
  dα = zeros(n)
  for i in 1:n
    dα[i] = ((1-y[i]*μ[i])^2+ζ[i,i])/(2*α[i])-0.5
  end
  dΘ = 0.5*(trace(invK*J*(invK*ζ-1))+dot(μ,invK*J*invK*μ))
  return [trace(dζ),mean(dμ),mean(dα),mean(dΘ)]
end

function SparseELBO(model)
  ELBO = 0.0
  ELBO += 0.5*(logdet(model.ζ)+logdet(model.invKmm))
  ELBO += -0.5*(sum(model.invKmm.*model.ζ)+dot(model.μ,model.invKmm*model.μ)) #trace replaced by sum
  ELBO += dot(model.y,model.κ*model.μ)
  ELBO += sum(0.25*log.(model.α) + log.(besselk.(0.5,sqrt.(model.α))))
  ζtilde = model.κ*model.ζ*transpose(model.κ)
  for i in 1:model.nSamples
    ELBO += 0.5/sqrt.(model.α[i])*(model.α[i]-(1-model.y[i]*dot(model.κ[i,:],model.μ))^2-(ζtilde[i,i]+model.Ktilde[i]))
  end
  return ELBO
end

function NoisySparseELBO(model,indices)
    n = length(indices)
    ELBO = 0.0
    ELBO += 0.5*(logdet(model.ζ)+logdet(model.invKmm))
    ELBO += -0.5*(sum(model.invKmm.*model.ζ)+dot(model.μ,model.invKmm*model.μ)) #trace replaced by sum
    ELBO += dot(model.y,model.κ*model.μ)
    ELBO += sum(0.25*log.(model.α[indices]) + log.(besselk.(0.5,sqrt.(model.α[indices]))))
    ζtilde = model.κ[indices,:]*model.ζ*transpose(model.κ[indices,:])
    for i in 1:n
      ELBO += 0.5/sqrt.(model.α[indices[i]])*(model.α[indices[i]]-(1-model.y[indices[i]]*dot(model.κ[indices[i],:],model.μ))^2-(ζtilde[i,i]+model.Ktilde[i]))
    end
    return ELBO
end

function ParametersPlotting(model::BSVM;option::String="All")
  if !model.Storing
    warn("Data was not saved during training, please rerun training with option Storing=true")
    return
  elseif isempty(model.StoredValues )
    warn("Model was not trained yet, please run a dataset before");
    return
  end
  figure("Evolution of model properties over time");
  iterations =  floor(Int64,1:size(model.evol_β,1))-1
  nsiterations = size(model.StoredValues,1);
  sparseiterations = floor(Int64,linspace(0,(nsiterations-1)*model.StoringFrequency,nsiterations))
  if option == "All"
    nFeatures = model.Autotuning ? 6 : 4;
    subplot(nFeatures÷2,2,1)
    plot(sparseiterations,model.StoredValues[:,1])
    ylabel("ELBO")
    subplot(nFeatures÷2,2,2)
    plot(sparseiterations,model.StoredValues[:,2])
    ylabel(L"Trace($\zeta$)")
    subplot(nFeatures÷2,2,3)
    plot(sparseiterations,model.StoredValues[:,3])
    ylabel(L"Mean($\mu$)")
    subplot(nFeatures÷2,2,4)
    plot(sparseiterations,model.StoredValues[:,4])
    ylabel(L"Mean($\alpha_i$)")
    if model.Autotuning
      subplot(nFeatures÷2,2,5)
      semilogy(sparseiterations,model.StoredValues[:,5])
      ylabel(model.NonLinear ? L"||\theta||" : L"\gamma")
      subplot(nFeatures÷2,2,6)
      semilogy(sparseiterations,model.StoredValues[:,6])
      ylabel(L"\rho_\theta")
    end
  elseif option == "dELBO"
    (nIterations,nFeatures) = size(model.StoredDerivativesELBO)
    if model.Sparse
        warn("ELBO derivatives not available yet for the Sparse Algorithm")
    end
    DerivativesLabels = ["d\zeta" "d\mu" "d\alpha" "d\theta"]
    for i in 1:4
      if i <= 3  || (i==4 && model.Autotuning)
        subplot(2,2,i)
        plot(sparseiterations,model.StoredDerivativesELBO[:,i])
        ylabel(DerivativesLabels[i])
      end
    end
elseif option == "Mu"
    plot(iterations,sqrt.(sumabs2(model.evol_β)))
    ylabel(L"Normalized $\mu$")
    xlabel("Iterations")
  elseif option == "ELBO"
    plot(iterations,model.StoredValues[:,2])
    ylabel("ELBO")
    xlabel("Iterations")
  else
    warn("Option not available, chose among those : All, dELBO, Mu, Autotuning, ELBO")
  end
  return;
end

function LinearPredict(X_test,model)
  return model.Intercept ? [ones(Float64,size(X_test,1)) X_test]*model.μ : X_test*model.μ
end

function LinearPredictProb(X_test,model)
  if model.Intercept
    X_test = [ones(Float64,size(X_test,1)) X_test]
  end
  n = size(X_test,1)
  predic = zeros(n)
  for i in 1:n
    predic[i] = cdf(Normal(),(dot(X_test[i,:],model.μ))/(dot(X_test[i,:],model.ζ*X_test[i,:])+1))
  end
  return predic
end

function NonLinearPredict(X_test,model)
  n = size(X_test,1)
  if model.top == 0
    model.top = model.invK*model.μ
  end
  k_star = CreateKernelMatrix(X_test,model.Kernel_function,X2=model.X)
  return k_star*model.top
end

function NonLinearPredictProb(X_test,model)
  n = size(X_test,1)
  if model.down == 0
    if model.invK == 0
      model.invK = Matrix(Symmetric(inv(CreateKernelMatrix(model.X,model.kernelfunction,model.Θ)+model.γ*eye(model.nSamples)),:U))
    end
    model.top = model.invK*model.μ
    model.down = -(model.invK+model.invK*model.ζ*model.invK)
  end
  ksize = model.nSamples
  predic = zeros(n)
  k_star = zeros(ksize)
  k_starstar = 0
  for i in 1:n
    for j in 1:ksize
      k_star[j] = model.Kernel_function(model.X[j,:],X_test[i,:])
    end
    k_starstar = model.Kernel_function(X_test[i,:],X_test[i,:])
    predic[i] = cdf(Normal(),(dot(k_star,model.top))/(k_starstar + dot(k_star,model.down*k_star) + 1))
  end
  predic
end

function SparsePredictProb(X_test,model)
  n = size(X_test,1)
  ksize = model.m
  if model.down == 0
    if model.top == 0
      model.top = model.invKmm*model.μ
    end
    model.down = model.invKmm*(eye(ksize)+model.ζ*model.invKmm)
    model.MatricesPrecomputed = true
  end
  predic = zeros(n)
  k_star = zeros(ksize)
  k_starstar = 0
  for i in 1:n
    for j in 1:ksize
      k_star[j] = model.Kernel_function(model.inducingPoints[j,:],X_test[i,:])
    end
    k_starstar = model.Kernel_function(X_test[i,:],X_test[i,:])

    predic[i] = cdf(Normal(),(dot(k_star,model.top))/(k_starstar - dot(k_star,model.down*k_star) + 1))
  end
  predic
end

function SparsePredict(X_test,model)
  n = size(X_test,1)
  if model.top == 0
    model.top = model.invKmm*model.μ
  end
  k_star = CreateKernelMatrix(X_test,model.Kernel_function,X2=model.inducingPoints)
  return k_star*model.top
end

end #End Module


text:
using LabView0mqJlExamples

fns = (;arr_amplitude, arr_length, arr_mean, arr_median, arr_rms, vec_amplify, vec_attenuate, vec_power, )


text:
using GISSMOReader
using Test

@testset "GISSMOReader.jl" begin
    # Write your tests here.
end


text:
struct ProbitModel <: GeneralisedLinearModel
    y::Vector
    X::Matrix
end

struct FittedProbitModel <: Fit
    y::Vector
    X::Matrix
    modelfit::FitOutput
end

function fit!(model::ProbitModel)
    X = model.X

    init_β = repeat([0.0], size(X, 2))

    f(β) = loglikelihood(model, β)
    mle_result = optimize(f, init_β, LBFGS())
    β = mle_result.minimizer
    # forgive me father for I have sinned
    # TODO
    return FitOutput(β, [0], [0], [0.0], [0 0; 0 0])
end

function fit(model::ProbitModel)
    FittedProbitModel(
        model.y,
        model.X,
        fit!(model))
end

function loglikelihood(model::ProbitModel, β)
    μ = model.X*β
    y = model.y
    ϕ = Normal()
    ll = sum(y .* log.(cdf.(ϕ, μ)) .+ (1 .- y) .* log.(1 .- cdf.(ϕ, μ)))
    return -ll
end

max_stars_repo_name: EdJeeOnGitHub/jeeves

max_issues_repo_name: EdJeeOnGitHub/jeeves

max_forks_repo_name: EdJeeOnGitHub/jeeves


text:
module BenchAppendBB

using BangBang: append!!
using Transducers
using BenchmarkTools: @benchmarkable, BenchmarkGroup

const SUITE = BenchmarkGroup()

SUITE["xf"] = @benchmarkable(
    append!!(Transducer(y for x in nothing if isodd(x) for y in 1:x), dest, 1:100),
    setup = (dest = Union{}[]),
)
SUITE["eduction"] = @benchmarkable(
    append!!(dest, eduction(y for x in 1:100 if isodd(x) for y in 1:x)),
    setup = (dest = Union{}[]),
)

end  # module
BenchAppendBB.SUITE


text:
using LazyBandedMatrices, BlockBandedMatrices, BlockArrays, StaticArrays, FillArrays, LazyArrays, ArrayLayouts, BandedMatrices, Test
import LazyBandedMatrices: BlockBroadcastArray, ApplyLayout, blockcolsupport, blockrowsupport, arguments, paddeddata, resizedata!
import LinearAlgebra: Adjoint, Transpose
import LazyArrays: PaddedArray

@testset "unitblocks" begin
    a = unitblocks(Base.OneTo(5))
    @test a == 1:5
    @test blockaxes(a,1) == Block.(1:5)

    a = unitblocks(2:5)
    @test a == 2:5
    @test blockaxes(a,1) == Block.(1:4)
end

@testset "BlockVcat" begin
    @testset "vec vcat" begin
        a = BlockVcat(1:5, 10:12, 14:15)
        @test @inferred(axes(a)) ≡ (blockedrange(SVector(5,3,2)),)
        @test @inferred(a[Block(1)]) ≡ 1:5
        @test a == [1:5; 10:12; 14:15]
        @test a[Block.(1:2)] ≡ BlockVcat(1:5, 10:12)
        @test blocklengths(axes(a[Block.(1:2)],1)) == [5,3]
        @test a[:] == a[1:size(a,1)] == a
        @test a[1:10] isa Vcat
        @test a[Block(1)[1:2]] ≡ 1:2
        @test a[3] ≡ 3

        @test copy(a) ≡ convert(AbstractArray{Int},a) ≡ convert(AbstractVector{Int},a) ≡ a
        @test AbstractArray{Float64}(a) == AbstractVector{Float64}(a) == convert(AbstractArray{Float64},a) == convert(AbstractVector{Float64},a) == a
        @test copy(a') ≡ a'
    end
    @testset "mat vcat" begin
        A = BlockVcat(randn(2,3), randn(3,3))
        @test axes(A,2) ≡ Base.OneTo(3)
        @test A[Block(1,1)] == A.arrays[1]
        @test A[Block.(1:2),Block(1)] == A
        @test A[Block.(1:2),Block(1)] isa typeof(A)
        @test blocklengths(axes(A[Block.(1:2),Block(1)],1)) == [2,3]
        @test blocklengths(axes(A[Block.(1:2),Block(1)],2)) == [3]

        @test A[Block.(1:2), Block.(1:1)] == A
        @test A[Block.(1:2), Block.(1:1)] isa BlockVcat
        @test A[Block.(1:2), 1:2] == A[:,1:2]
        @test A[Block.(1:2), 1:2] isa BlockVcat

        @test convert(AbstractArray{Float64},A) == convert(AbstractMatrix{Float64},A) == A
        @test copy(A) == AbstractArray{Float64}(A) == AbstractMatrix{Float64}(A) == A
        @test copy(A') == A'
    end
    @testset "block vec vcat" begin
        a = PseudoBlockArray(1:5, SVector(1,3))
        b = PseudoBlockArray(2:6, SVector(1,3))

        c = BlockVcat(a,b)
        @test c == [a; b]
        @test c[Block(2)] == a[Block(2)]
        @test c[Block(3)] == b[Block(1)]
        @test c[Block.(2:3)] == [a[Block(2)]; b[Block(1)]]
        @test c[Block.(2:3),Block(1)] == reshape([a[Block(2)]; b[Block(1)]], 4, 1)
        @test copy(c) ≡ convert(AbstractArray{Int},c) ≡ convert(AbstractVector{Int},c) ≡ c
        @test AbstractArray{Float64}(c) == AbstractVector{Float64}(c) == c
        @test copy(c') ≡ c'

        A = BlockVcat(a', b')
        @test axes(A,1) ≡ blockedrange(SVector(1,1))
        @test axes(A,2) ≡ axes(a,1)
        @test A[Block(1,1)] == a[Block(1)]'
        @test A[Block(2,2)] == b[Block(2)]'
        @test A == [a'; b']
        @test convert(AbstractArray{Int},A) ≡ convert(AbstractMatrix{Int},A) ≡ A
        @test copy(A) == AbstractArray{Float64}(A) == AbstractMatrix{Float64}(A) == A
        @test copy(A') == A'
    end

    @testset "block mat vcat" begin
        A = PseudoBlockArray(randn(3,2), [1,2], [1,1])
        B = PseudoBlockArray(randn(4,2), [3,1], [1,1])
        V = BlockVcat(A, B)
        @test V == [A; B]
        @test V[Block(3,1)] == B[Block(1,1)]
        @test convert(AbstractArray{Float64},V) ≡ convert(AbstractMatrix{Float64},V) ≡ V
        @test copy(V) == AbstractArray{Float64}(V) == AbstractMatrix{Float64}(V) == V
        @test copy(A') == A'

        @test V[Block.(2:3), Block.(1:2)] == [A[Block(2),Block.(1:2)]; B[Block(1),Block.(1:2)]]
    end

    @testset "triangle recurrence" begin
        N = 1000
        a = b = c = 0.0
        n = mortar(Fill.(Base.OneTo(N),Base.OneTo(N)))
        k = mortar(Base.OneTo.(Base.OneTo(N)))
        dat = BlockVcat(
            ((k .+ (c-1)) .* ( k .- n .- 1 ) ./ (2k .+ (b+c-1)))',
            (k .* (k .- n .- a) ./ (2k .+ (b+c-1)))'
            )
        @test dat[:,Block.(1:3)] isa BlockVcat
    end

    @testset "PseudoBlockArray" begin
        a = PseudoBlockArray(Vcat(randn(3), 1:3), [3,3])
        @test MemoryLayout(a) isa LazyArrays.ApplyLayout{typeof(vcat)}
        @test LazyArrays.arguments(a) == LazyArrays.arguments(a.blocks)

        b = PseudoBlockArray(Vcat(randn(3), Zeros(3)), [3,3])
        @test paddeddata(view(b, 1:4)) == paddeddata(view(b, Base.OneTo(4))) == b[1:3]

        c = PseudoBlockArray(cache(Zeros(6)), 1:3);
        c[2] = 2
        @test blocksize(paddeddata(c)) == (2,)
        @test paddeddata(c)[Block(2)] == [2.0,0.0]
        resizedata!(c,4);
        @test blocksize(paddeddata(c)) == (3,)

        dat = randn(3,3)
        A = PseudoBlockArray(PaddedArray(dat, 6,6), 1:3, 1:3)
        @test paddeddata(A) == dat
    end

    @testset "blockcol/rowsupport" begin
        B = BlockBandedMatrix(randn(10,10),1:4,1:4,(1,0))
        V = BlockVcat(B, B)
        @test blockrowsupport(V, Block(3)) == blockrowsupport(V, Block(7)) == Block.(2:3)
        @test blockcolsupport(V, Block(1)) == Block.(1:6)
        @test blockbandwidths(V) == (5,0)
        @test isblockbanded(V)
    end
end

@testset "BlockHcat" begin
    @testset "vec hcat" begin
        a = BlockHcat(1:5, 10:14)
        @test axes(a,2) ≡ blockedrange(Ones{Int}(2))
        @test a[Block(1,1)] ≡ 1:5
        @test a == [1:5 10:14]
        @test_broken a[:,Block.(1:2)] ≡ BlockHcat(1:5, 10:14)
        @test a[:] == a[1:length(a)] == vec(a)
        @test copy(a) ≡ convert(AbstractArray{Int},a) ≡ convert(AbstractMatrix{Int},a) ≡ a
        @test AbstractArray{Float64}(a) == AbstractMatrix{Float64}(a) == convert(AbstractArray{Float64},a) == convert(AbstractMatrix{Float64},a) == a
        @test copy(a') == a'
    end

    @testset "mat hcat" begin
        A = BlockHcat(randn(3,2), randn(3,3))
        @test axes(A,1) ≡ Base.OneTo(3)
        @test A[Block(1,1)] == A.arrays[1]
        @test A[Block(1),Block.(1:2)] == A
        @test A[Block(1),Block.(1:2)] isa BlockHcat
        @test convert(AbstractArray{Float64},A) ≡ convert(AbstractMatrix{Float64},A) ≡ A
        @test copy(A) == AbstractArray{Float64}(A) == AbstractMatrix{Float64}(A) == A
        @test copy(A') == copy(Adjoint(A)) == A'
    end

    @testset "block vec hcat" begin
        a = PseudoBlockArray(1:4, SVector(1,3))
        b = PseudoBlockArray(2:5, SVector(1,3))
        A = BlockHcat(a, b)
        @test axes(A,2) ≡ blockedrange(Ones{Int}(2))
        @test axes(A,1) ≡ axes(a,1)
        @test A[Block(1,1)] == a[Block(1)]
        @test A[Block(2,2)] == b[Block(2)]
        @test A[Block(2), Block.(1:2)] == [a[Block(2)] b[Block(2)]]
        @test A[Block.(1:2),Block.(2)] == reshape(b,4,1)
        @test A[Block.(1:2), Block.(1:1)] == reshape(a,4,1)

        @test A == [a b]
        @test convert(AbstractArray{Int},A) ≡ convert(AbstractMatrix{Int},A) ≡ A
        @test copy(A) == AbstractArray{Float64}(A) == AbstractMatrix{Float64}(A) == A
        @test copy(A') == A'
    end

    @testset "block mat hcat" begin
        A = PseudoBlockArray(randn(3,2), [1,2], [1,1])
        B = PseudoBlockArray(randn(3,3), [1,2], [2,1])
        H = BlockHcat(A, B)
        @test H == [A B]
        @test H[Block(1,3)] == B[Block(1,1)]

        @test H[Block(2),Block.(2:3)] == H[Block.(2:2),Block.(2:3)]  == [A[Block(2,2)] B[Block(2,1)]]
        @test H[Block.(2:2),Block(3)] == H[Block(2,3)]
        @test convert(AbstractArray{Float64},H) ≡ convert(AbstractMatrix{Float64},H) ≡ H
        @test copy(H) == AbstractArray{Float64}(H) == AbstractMatrix{Float64}(H) == H
        @test copy(H') == H'
    end

    @testset "triangle recurrence" begin
        N = 1_000
        a = b = c = 0.0
        n = mortar(BroadcastArray(Fill,Base.OneTo(N),Base.OneTo(N)))
        k = mortar(BroadcastArray(Base.OneTo,Base.OneTo(N)))
        A = BlockHcat(
            BroadcastVector((n,k,bc1,abc) -> (n + k +  bc1) / (2n + abc), n, k, b+c-1, a+b+c),
            BroadcastVector((n,k,abc) -> (n + k +  abc) / (2n + abc), n, k, a+b+c)
            )
        dest = PseudoBlockArray{Float64}(undef, axes(A))
        @test copyto!(dest, A) == A;
        @test @allocated(copyto!(dest, A)) ≤ 2800
        # dest = BlockArray{Float64}(undef, axes(A))
        # @time copyto!(dest, A);

        dest = PseudoBlockArray{Float64}(undef, axes(A'))
        @test (A')[Block(2,3)] == A[Block(3,2)]'
        @test copyto!(dest, A') == A'
        @test @allocated(copyto!(dest, A')) ≤ 2600

        Rx = BlockBandedMatrices._BandedBlockBandedMatrix(A', axes(k,1), (0,1), (0,0))
        dest = BandedBlockBandedMatrix{Float64}(undef, axes(Rx), (0,1), (0,0))
        @test copyto!(dest, Rx) == BandedBlockBandedMatrix(Rx)

        Vx = view(Rx, Block.(1:N), Block.(1:N))
        @test MemoryLayout(Vx) isa BlockBandedMatrices.BandedBlockBandedColumns

        V = view(A, Block.(1:5),:)
        @test MemoryLayout(V) isa ApplyLayout{typeof(hcat)}
        HV = A[Block.(1:5),:]
        @test HV == V
        # TODO: Fast materialization
    end

    # @testset "Block-mat hcat" begin
    #     A = BlockHcat(PseudoBlockArray(randn(6,4), [4,2], [2,2]), PseudoBlockArray(randn(6,3), [4,2], [2,1]))
    #     A[:,Block.(1:2)]
    # end

    @testset "adjtrans" begin
        A = BlockHcat(randn(3,2), randn(3,3))
        @test A' == transpose(A) == BlockArray(A)'
        @test A' isa BlockVcat
        @test transpose(A) isa BlockVcat
        @test A'' == A
        @test A'' isa BlockHcat
    end

    @testset "blockcol/rowsupport" begin
        B = BlockBandedMatrix(randn(10,10),1:4,1:4,(1,0))
        H = BlockHcat(B, B)
        @test blockcolsupport(H, Block(2)) == blockcolsupport(H, Block(6)) == Block.(2:3)
        @test blockrowsupport(H, Block(1)) == Block.(1:5)
        @test blockbandwidths(H) == (1,4)
        @test isblockbanded(H)
        @test copyto!(similar(H), H) == copy(H) == H
    end

    @testset "Eye BlockHcat" begin
        B = BlockBandedMatrix(randn(10,10),1:4,1:4,(1,0))
        H = BlockHcat(Eye((axes(B,1),))[:,Block(1)], B)
        @test MemoryLayout(H) isa LazyBandedMatrices.ApplyBlockBandedLayout{typeof(hcat)}
        @test H[Block.(1:4), Block.(1:5)] == H == copy(H)
    end
end

@testset "BlockHvcat" begin
    A = randn(2,2)
    B = randn(2,3)
    C = randn(3,2)
    D = randn(3,3)

    H = BlockHvcat(2, A, B, C, D)
    @test H == [A B; C D]

    @test convert(AbstractArray{Float64},H) ≡ convert(AbstractMatrix{Float64},H) ≡ H
    @test copy(H) == AbstractArray{Float64}(H) == AbstractMatrix{Float64}(H) == H
    @test copy(H') == H'
end

@testset "Interlace" begin
    @testset "vcat" begin
        N = 1000
        a = 1:N
        b = 11:10+N
        a, b = PseudoBlockArray(a,Ones{Int}(length(a))), PseudoBlockArray(b,Ones{Int}(length(b)))
        A = BlockBroadcastArray(vcat, a, b)
        if VERSION < v"1.7-"
            @test axes(A,1) isa BlockedUnitRange{StepRange{Int,Int}}
        else
            @test axes(A,1) isa BlockedUnitRange{StepRangeLen{Int,Int,Int,Int}}
        end
        @test @allocated(axes(A)) ≤ 50
        @test A[Block(1)] == PseudoBlockArray(A)[Block(1)] == [A[1],A[2]] == [1,11]
        @test A[Block(N)] == PseudoBlockArray(A)[Block(N)] == [1000,1010]
        @test convert(AbstractArray{Int},A) ≡ convert(AbstractVector{Int},A) ≡ A
        @test copy(A) == AbstractArray{Float64}(A) == AbstractVector{Float64}(A) == convert(AbstractArray{Float64},A) == convert(AbstractVector{Float64},A) == A
        @test copy(A') == A'

        @testset "padded" begin
            a = Vcat(randn(3), Zeros(10))
            b = Vcat(randn(4), Zeros(9))
            C = BlockBroadcastArray(vcat,unitblocks(a),unitblocks(a))
            @test C[1:2:end] == a
            @test C[2:2:end] == a

            # differening data sizes not supported yet
            @test_broken paddeddata(BlockBroadcastArray(vcat,unitblocks(a),unitblocks(b)))
        end

        @testset "resize!" begin
            N = 10
            a, b = PseudoBlockArray(randn(N),Ones{Int}(N)), PseudoBlockArray(randn(N),Ones{Int}(randn(N)))
            A = BlockBroadcastArray(vcat, a, b)
            Ã = resize!(A, Block(3))
            @test Ã == A[1:6]
            @test_throws BoundsError A[7]
            @test !isassigned(A,7)
        end
    end
    @testset "hcat" begin
        N = 1000
        a = 1:N
        b = 11:10+N
        a, b = PseudoBlockArray(a,Ones{Int}(length(a))), PseudoBlockArray(b,Ones{Int}(length(b)))
        A = BlockBroadcastArray(hcat, a', b')
        if VERSION < v"1.7-"
            @test axes(A,2) isa BlockedUnitRange{StepRange{Int,Int}}
        else
            @test axes(A,2) isa BlockedUnitRange{StepRangeLen{Int,Int,Int,Int}}
        end
        @test @allocated(axes(A)) ≤ 50
        @test A[Block(1,1)] == PseudoBlockArray(A)[Block(1,1)] == [1 11]
        @test A[Block(1,N)] == PseudoBlockArray(A)[Block(1,N)] == [1000 1010]
        @test convert(AbstractArray{Int},A) ≡ convert(AbstractMatrix{Int},A) ≡ A
        @test copy(A) == AbstractArray{Float64}(A) == AbstractMatrix{Float64}(A) == A
        @test copy(A') == A'
    end
    @testset  "hvcat" begin
        a = unitblocks(randn(2,3))
        b = unitblocks(randn(2,3))
        c = unitblocks(randn(2,3))
        d = unitblocks(randn(2,3))
        e = unitblocks(randn(2,3))
        f = unitblocks(randn(2,3))

        A = BlockBroadcastArray(hvcat, 2, a, b, c, d, e, f)
        @test MemoryLayout(A) isa LazyLayout
        @test blocksize(A) == (2,3)
        @test A[Block(1,1)] == [a[1] b[1]; c[1] d[1]; e[1] f[1]]
        @test A[1,1] == a[1,1]

        @test convert(AbstractArray{Float64},A) ≡ convert(AbstractMatrix{Float64},A) ≡ A
        @test copy(A) == AbstractArray{Float64}(A) == AbstractMatrix{Float64}(A) == A
        @test convert(AbstractArray{BigFloat}, A) == convert(AbstractMatrix{BigFloat}, A) == A
        @test copy(A') == A'

        @testset "Banded" begin
            a = unitblocks(brand(5,4,1,2))
            z = Zeros(axes(a))
            @test blockbandwidths(a) == (1,2)
            @test subblockbandwidths(a) == (0,0)
            A = BlockBroadcastArray(hvcat, 2, a, z, z, a)
            @test blockbandwidths(A) == (1,2)
            @test subblockbandwidths(A) == (0,0)
            @test A[Block.(1:2),Block.(1:2)] isa BlockSkylineMatrix

            @test blockcolsupport(A, Block(2)) == Block.(1:3)
            @test blockrowsupport(A, Block(3)) == Block.(2:4)

            V = view(A, Block.(1:3),Block.(1:3))
            @test MemoryLayout(V) isa LazyBandedMatrices.BlockBandedInterlaceLayout
            @test arguments(V) == (2,a[1:3,1:3],z[1:3,1:3],z[1:3,1:3],a[1:3,1:3])
        end
    end

    @testset "Diagonal" begin
        A = unitblocks(brand(5,5,1,2))
        B = unitblocks(brand(5,5,2,1))
        C = BlockBroadcastArray{Float64}(Diagonal, A, B)
        @test blockisequal(axes(C),(blockedrange(Fill(2,5)),blockedrange(Fill(2,5))))
        for k = 1:5, j = 1:5
            @test C[Block(k,j)] == Diagonal([A[k,j],B[k,j]])
        end
        @test blockbandwidths(C) == (2,2)
        @test subblockbandwidths(C) == (0,0)
    end
end


text:
import jInv.ForwardShare.getSensMat
function getSensMat(m::Vector,pFor::PointCloudParam)
	return pFor.Jacobian;
end

export getSensMatVec
function getSensMatVec(v::Vector,m::Vector,pFor::PointCloudParam)
	println("Im in getSensMatVec")
	S = pFor.Jacobian;

	d = S'*v[:];

end

export getSensTMatVec
function getSensTMatVec(v::Vector,m::Vector,pFor::PointCloudParam)
	JtV = 0;
	S = pFor.Jacobian;
	println("v size:",size(v)); println("m size:",size(m)); println("J size:",size(S));
	JtV = S'*v[:];
return JtV;
end

max_stars_repo_name: BGUCompSci/ShapeReconstructionPaLS

max_issues_repo_name: BGUCompSci/ShapeReconstructionPaLS

max_forks_repo_name: BGUCompSci/ShapeReconstructionPaLS


text:
using Distributed

@everywhere begin
    include("base_script.jl")
end


text:
using DataDeps

register(DataDep(
  "abalone",
  "https://archive.ics.uci.edu/ml/datasets/Abalone",
  "https://archive.ics.uci.edu/ml/machine-learning-databases/abalone/abalone.data",
  "de37cdcdcaaa50c309d514f248f7c2302a5f1f88c168905eba23fe2fbc78449f",
  post_fetch_method=(path -> begin
    UCIData.process_dataset(path,
        target_index=9,
        feature_indices=1:8,
        categoric_indices=1,
    )
  end),
))

max_stars_repo_name: JackDunnNZ/uci-data

max_issues_repo_name: JackDunnNZ/uci-data

max_forks_repo_name: JackDunnNZ/uci-data


text:
"""
"""
struct LocalToGlobalArray{T,M,N,L,V} <: AbstractArray{Array{T,M},N}
  lid_to_gid::L
  gid_to_val::V

  @doc """
  """
  function LocalToGlobalArray(
    lid_to_gid::AbstractArray{<:AbstractArray},
    gid_to_val::AbstractArray)

    L = typeof(lid_to_gid)
    V = typeof(gid_to_val)
    T = eltype(V)
    M = ndims(eltype(L))
    N = ndims(L)
    new{T,M,N,L,V}(lid_to_gid,gid_to_val)
  end
end

size(a::LocalToGlobalArray) = size(a.lid_to_gid)

function IndexStyle(::Type{LocalToGlobalArray{T,M,N,L,V}}) where {T,M,N,L,V}
  IndexStyle(L)
end

function array_cache(a::LocalToGlobalArray)
  gids = testitem(a.lid_to_gid)
  T = eltype(a.gid_to_val)
  r = zeros(T,size(gids))
  c = CachedArray(r)
  cl = array_cache(a.lid_to_gid)
  (cl,c)
end

function getindex!(cache,a::LocalToGlobalArray,i::Integer...)
  (cl,c) = cache
  gids = getindex!(cl,a.lid_to_gid,i...)
  setsize!(c,size(gids))
  r = c.array
  for i in eachindex(gids)
    @inbounds r[i] = a.gid_to_val[gids[i]]
  end
  r
end

function getindex(a::LocalToGlobalArray,i::Integer...)
  cache = array_cache(a)
  getindex!(cache,a,i...)
end

function reindex(a::LocalToGlobalArray,b::AbstractArray)
  lid_to_gid = reindex(a.lid_to_gid,b)
  LocalToGlobalArray(lid_to_gid,a.gid_to_val)
end


text:
function FCurlNon3Vec!(FuHat,v1CG,v2CG,wCG,wCCG,CG,Param)
OP=CG.OrdPoly+1
NF=Param.Grid.NumFaces
nz=Param.Grid.nz
uPos=Param.uPos
vPos=Param.vPos
wPos = Param.wPos
JC = Param.cache.JC

dXdxIC11 = Param.dXdxIC11
dXdxIC21 = Param.dXdxIC21
dXdxIC31 = Param.dXdxIC31
dXdxIC12 = Param.dXdxIC12
dXdxIC22 = Param.dXdxIC22
dXdxIC32 = Param.dXdxIC32
dXdxIC33 = Param.dXdxIC33

dXdxIC = Param.cache.dXdxIC
dXdxIF = Param.cache.dXdxIF

vHat = Param.CacheC1
Vort3 = Param.CacheC2
Temp = Param.CacheC3
DZvHat3 = Param.CacheC4

wHat = Param.CacheF1
DXwHat11 = Param.CacheF2
DYwHat21 = Param.CacheF3
DZuuHat31 = Param.CacheF4
DZwwHat31 = Param.CacheC5

DXwHat12 = Param.CacheF2
DYwHat22 = Param.CacheF3
DZuuHat32 = Param.CacheF4
DZwwHat32 = Param.CacheC5

# Fu(1)
# -vv\left(DX(a^{12}u-a^{11}v)+(a^{22}u-a^{21}v)DY^T
# +\frac{d(a^{32}u-a^{31}v)}{dz}\right)
#
# ww(DXa^{11}w+a^{21}wDY^T+\frac{da^{31}w}{dz}-\frac{da^{33}u}{dz})

# Fu(2)
# uu\left(DX(a^{12}u-a^{11}v)+(a^{22}u-a^{21}v)DY^T
# +\frac{d(a^{32}u-a^{31}v)}{dz}\right)
#
# ww(DXa^{12}w+a^{22}wDY^T+\frac{da^{32}w}{dz}-\frac{d a^{33}v}{dz})

# Fw
# -uu((DXa^{11}w+a^{21}wDY^T+\frac{da^{31}w}{dz}-\frac{da^{33}u}{dz})
# -vv*((DXa^{12}w+a^{22}wDY^T+\frac{da^{32}w}{dz}-\frac{da^{33}v}{dz})

vHat = dXdxIC12.*v1CG .- dXdxIC11.*v2CG
mul!(reshape(Vort3,OP,OP*NF*nz),CG.DS,reshape(vHat,OP,OP*nz*NF))
vHat = dXdxIC22.*v1CG .- dXdxIC21.*v2CG
mul!(reshape(PermutedDimsArray(Temp,(2,1,3,4)),OP,OP*NF*nz),CG.DS,reshape(PermutedDimsArray(vHat,(2,1,3,4)),OP,OP*nz*NF))
Vort3 .= Vort3 .+ Temp
vHat = dXdxIC32.*v1CG .- dXdxIC31.*v2CG
if nz>1
  @views Vort3[:,:,:,1] .+= 0.5 .* (vHat[:,:,:,2] .- vHat[:,:,:,1])
  @views Vort3[:,:,:,2:nz-1] .+= 0.25 .* (vHat[:,:,:,3:nz] .- vHat[:,:,:,1:nz-2])
  @views Vort3[:,:,:,nz] .+= 0.5 .* (vHat[:,:,:,nz] .- vHat[:,:,:,nz-1])
end
if Param.Coriolis
  str = Param.CoriolisType
  if str == "Sphere"
      lat = Param.cache.lat
      JC = Param.cache.JC
      Omega = Param.Omega
      @views Vort3 .= Vort3 .- reshape((2.0*Omega).* sin.(
        repmat(reshape(lat[:,:,:],OP*OP*NF,1),1,nz))
        ,OP,OP,NF,nz).* JC
  elseif str == "Beta-Plane"
      J = Param.cache.J
      X = Param.cache.X
      beta0 = Param.beta0
      f0 = Param.f0
      y0 = Param.y0
      @views Vort3 .= Vort3 .- reshape((f0+beta0 .* (
        reshape(abs.(X[:,:,2,:,:]),OP,OP,NF,nz) .- y0)).*
        J[:,:,:,:],OP*OP*NF,nz)
  end
end
@views FuHat[:,:,:,:,uPos] .-= Vort3.*v2CG
@views FuHat[:,:,:,:,vPos] .+= Vort3.*v1CG

# DXa^{11}w
@views wHat .= dXdxIF[:,:,:,:,1,1].*wCG
mul!(reshape(DXwHat11,OP,OP*NF*(nz+1)),CG.DS,reshape(wHat,OP,OP*(nz+1)*NF))
# a^{21}wDY^T
@views wHat .= dXdxIF[:,:,:,:,2,1].*wCG
mul!(reshape(PermutedDimsArray(DYwHat21,(2,1,3,4)),OP,OP*NF*(nz+1)),CG.DS,reshape(PermutedDimsArray(wHat,(2,1,3,4)),OP,OP*(nz+1)*NF))
# \frac{da^{31}w}{dz}-\frac{da^{33}u}{dz}
vHat .= dXdxIC33.*v1CG
if nz>1
  @views DZuuHat31[:,:,:,1] .= 0.5 .* (vHat[:,:,:,2] .- vHat[:,:,:,1])
  @views DZuuHat31[:,:,:,2:nz] .= 0.5 .* (vHat[:,:,:,2:nz] .- vHat[:,:,:,1:nz-1])
  @views DZuuHat31[:,:,:,nz+1] .= 0.5 .* (vHat[:,:,:,nz] .- vHat[:,:,:,nz-1])
end
@views wHat .= dXdxIF[:,:,:,:,3,2] .* wCG
@views DZwwHat31 .= 0.5 .* (wHat[:,:,:,2:nz+1] .- wHat[:,:,:,1:nz])
@views FuHat[:,:,:,:,uPos] .+=
  DZwwHat31 .* wCCG+
  0.5.*((DXwHat11[:,:,:,1:nz] .+ DYwHat21[:,:,:,1:nz] .-
  DZuuHat31[:,:,:,1:nz]).*wCG[:,:,:,1:nz] .+
  (DXwHat11[:,:,:,2:nz+1] .+ DYwHat21[:,:,:,2:nz+1] .-
  DZuuHat31[:,:,:,2:nz+1]).*wCG[:,:,:,2:nz+1])
@views FuHat[:,:,:,1:nz-1,wPos] .+=
  (-DXwHat11[:,:,:,2:nz] .- DYwHat21[:,:,:,2:nz] .+ DZuuHat31[:,:,:,2:nz]) .*
    (0.5 .* (v1CG[:,:,:,1:nz-1] .+ v1CG[:,:,:,2:nz])) .-
    0.5 .* (DZwwHat31[:,:,:,1:nz-1] .+ DZwwHat31[:,:,:,2:nz])

# DXa^{12}w
@views wHat .= dXdxIF[:,:,:,:,1,2].*wCG
mul!(reshape(DXwHat12,OP,OP*NF*(nz+1)),CG.DS,reshape(wHat,OP,OP*(nz+1)*NF))
# a^{22}wDY^T
@views wHat .= dXdxIF[:,:,:,:,2,2].*wCG
mul!(reshape(PermutedDimsArray(DYwHat22,(2,1,3,4)),OP,OP*NF*(nz+1)),CG.DS,reshape(PermutedDimsArray(wHat,(2,1,3,4)),OP,OP*(nz+1)*NF))
# \frac{da^{32}w}{dz}-\frac{da^{33}v}{dz}
vHat .= dXdxIC33 .* v2CG
if nz>1
  @views DZuuHat32[:,:,:,1] .= 0.5 .* (vHat[:,:,:,2] .- vHat[:,:,:,1])
  @views DZuuHat32[:,:,:,2:nz] .= 0.5 .* (vHat[:,:,:,2:nz] .- vHat[:,:,:,1:nz-1])
  @views DZuuHat32[:,:,:,nz+1] .= 0.5 .* (vHat[:,:,:,nz] .- vHat[:,:,:,nz-1])
end
@views wHat .= dXdxIF[:,:,:,:,3,2] .* wCG
@views DZwwHat32 .= 0.5 .* (wHat[:,:,:,2:nz+1] .- wHat[:,:,:,1:nz])

@views FuHat[:,:,:,:,vPos] .+=
  (DZwwHat32).*wCCG .+
  0.5 .* ((DXwHat12[:,:,:,1:nz]  .+ DYwHat22[:,:,:,1:nz] .-
    DZuuHat32[:,:,:,1:nz]).*wCG[:,:,:,1:nz] .+
  (DXwHat12[:,:,:,2:nz+1] .+ DYwHat22[:,:,:,2:nz+1] .-
    DZuuHat32[:,:,:,2:nz+1]).*wCG[:,:,:,2:nz+1])

DZwwHat32 .= DZwwHat32 .* v2CG
@views FuHat[:,:,:,1:nz-1,wPos] .+=
    (-DXwHat12[:,:,:,2:nz] .- DYwHat22[:,:,:,2:nz] .+ DZuuHat32[:,:,:,2:nz]) .*
    (0.5 .* (v2CG[:,:,:,1:nz-1] .+ v2CG[:,:,:,2:nz])) .-
    0.5 .* (DZwwHat32[:,:,:,1:nz-1] .+ DZwwHat32[:,:,:,2:nz])
end

function FCurlNon3Vec(v1CG,v2CG,wCG,wCCG,CG,Param)
OP=CG.OrdPoly+1
NF=Param.Grid.NumFaces
nz=Param.Grid.nz
dXdxIC = Param.cache.dXdxIC
dXdxIF = Param.cache.dXdxIF
# Fu(1)
# -vv\left(DX(a^{12}u-a^{11}v)+(a^{22}u-a^{21}v)DY^T
# +\frac{d(a^{32}u-a^{31}v)}{dz}\right)
#
# ww(DXa^{11}w+a^{21}wDY^T+\frac{da^{31}w}{dz}-\frac{da^{33}u}{dz})

# Fu(2)
# uu\left(DX(a^{12}u-a^{11}v)+(a^{22}u-a^{21}v)DY^T
# +\frac{d(a^{32}u-a^{31}v)}{dz}\right)
#
# ww(DXa^{12}w+a^{22}wDY^T+\frac{da^{32}w}{dz}-\frac{d a^{33}v}{dz})

# Fw
# -uu((DXa^{11}w+a^{21}wDY^T+\frac{da^{31}w}{dz}-\frac{da^{33}u}{dz})
# -vv*((DXa^{12}w+a^{22}wDY^T+\frac{da^{32}w}{dz}-\frac{da^{33}v}{dz})

vHat1 = dXdxIC[:,:,:,:,1,2].*v1CG .- dXdxIC[:,:,:,:,1,1].*v2CG
vHat2 = dXdxIC[:,:,:,:,2,2].*v1CG .- dXdxIC[:,:,:,:,2,1].*v2CG
vHat3 = dXdxIC[:,:,:,:,3,2].*v1CG .- dXdxIC[:,:,:,:,3,1].*v2CG

DXvHat1=reshape(
  CG.DS*reshape(vHat1,OP,OP*nz*NF),
  OP,OP,NF,nz)
DYvHat2= permute(
  reshape(
  CG.DS*reshape(
  permute(vHat2
  ,[2,1,3,4])
  ,OP,OP*NF*nz)
  ,OP,OP,NF,nz)
  ,[2,1,3,4])

DZvHat3=zeros(OP,OP,NF,nz)
if nz>1
  DZvHat3[:,:,:,1]=0.5*(vHat3[:,:,:,2]-vHat3[:,:,:,1])
  DZvHat3[:,:,:,2:nz-1]=0.25*(vHat3[:,:,:,3:nz]-vHat3[:,:,:,1:nz-2])
  DZvHat3[:,:,:,nz]=0.5*(vHat3[:,:,:,nz]-vHat3[:,:,:,nz-1])
end

# DXa^{11}w
DXwHat11=reshape(CG.DS*reshape(
  dXdxIF[:,:,:,:,1,1].*wCG
  ,OP,OP*NF*(nz+1))
  ,OP,OP,NF,nz+1)
# DXa^{12}w
DXwHat12=reshape(CG.DS*reshape(
  dXdxIF[:,:,:,:,1,2].*wCG
  ,OP,OP*NF*(nz+1))
  ,OP,OP,NF,nz+1)
# a^{21}wDY^T
DYwHat21=  permute(
  reshape(
  CG.DS*reshape(
  permute(
  reshape(
  dXdxIF[:,:,:,:,2,1].*wCG
  ,OP,OP,NF,nz+1)
  ,[2,1,3,4])
  ,OP,OP*NF*(nz+1))
  ,OP,OP,NF,nz+1)
  ,[2,1,3,4])
# a^{22}wDY^T
DYwHat22= permute(
  reshape(
  CG.DS*reshape(
  permute(
  reshape(
  dXdxIF[:,:,:,:,2,2].*wCG
  ,OP,OP,NF,nz+1)
  ,[2,1,3,4])
  ,OP,OP*NF*(nz+1))
  ,OP,OP,NF,nz+1)
  ,[2,1,3,4])

# \frac{da^{31}w}{dz}-\frac{da^{33}u}{dz}
#uHat1=reshape(Param.dXdxI(:,:,:,:,3,3),OP*OP*NF,nz).*v1CG
uHat1=dXdxIC[:,:,:,:,3,3].*v1CG
DZuuHat31=zeros(OP,OP,NF,nz+1)
if nz>1
  DZuuHat31[:,:,:,1]=0.5*(uHat1[:,:,:,2]-uHat1[:,:,:,1])
  DZuuHat31[:,:,:,2:nz]=0.5*(uHat1[:,:,:,2:nz]-uHat1[:,:,:,1:nz-1])
  DZuuHat31[:,:,:,nz+1]=0.5*(uHat1[:,:,:,nz]-uHat1[:,:,:,nz-1])
end

wwHat31=dXdxIF[:,:,:,:,3,1].*wCG
DZwwHat31=0.5*(wwHat31[:,:,:,2:nz+1]-wwHat31[:,:,:,1:nz])
# \frac{da^{32}w}{dz}-\frac{da^{33}v}{dz}
#uHat2=reshape(Param.dXdxI(:,:,:,:,3,3),OP*OP*NF,nz).*v2CG
uHat2=dXdxIC[:,:,:,:,3,3].*v2CG
DZuuHat32=zeros(OP,OP,NF,nz+1)
if nz>1
  DZuuHat32[:,:,:,1]=0.5*(uHat2[:,:,:,2]-uHat2[:,:,:,1])
  DZuuHat32[:,:,:,2:nz]=0.5*(uHat2[:,:,:,2:nz]-uHat2[:,:,:,1:nz-1])
  DZuuHat32[:,:,:,nz+1]=0.5*(uHat2[:,:,:,nz]-uHat2[:,:,:,nz-1])
end
wwHat32=dXdxIF[:,:,:,:,3,2].*wCG
DZwwHat32=0.5*(wwHat32[:,:,:,2:nz+1]-wwHat32[:,:,:,1:nz])

FuHat=zeros(OP,OP,NF,nz,3)
#FuHat[:,:,:,1]=reshape(-(DXvHat1+DYvHat2+DZvHat3).*v2CG+(DXwHat11+DYwHat21+DZwHat31).*wC
#  ,OP*OP,NF,nz)
# Fu(1)
# -vv\left(DX(a^{12}u-a^{11}v)+(a^{22}u-a^{21}v)DY^T
# +\frac{d(a^{32}u-a^{31}v)}{dz}\right)
#
# ww(DXa^{11}w+a^{21}wDY^T+\frac{da^{31}w}{dz}-\frac{da^{33}u}{dz})
Vort3=DXvHat1+DYvHat2+DZvHat3
if Param.Coriolis
  str = Param.CoriolisType
  if str == "Sphere"
      lat = Param.cache.lat
      JC = Param.cache.JC
      Omega = Param.Omega
      Vort3=Vort3-reshape(2.0*Omega*sin.(
        repmat(reshape(lat[:,:,:],OP*OP*NF,1),1,nz))
        ,OP,OP,NF,nz).*
        JC
  elseif str == "Beta-Plane"
      J = Param.cache.J
      X = Param.cache.X
      beta0 = Param.beta0
      f0 = Param.f0
      y0 = Param.y0
      Vort3=Vort3-reshape((f0+beta0*(
        reshape(abs.(X[:,:,2,:,:]),OP,OP,NF,nz)-y0)).*
        J[:,:,:,:],OP*OP*NF,nz)
  end
end
FuHat[:,:,:,:,1]= -Vort3.*v2CG+
  (DZwwHat31).*wCCG+
  0.5*((DXwHat11[:,:,:,1:nz]  +DYwHat21[:,:,:,1:nz]-
  DZuuHat31[:,:,:,1:nz]).*wCG[:,:,:,1:nz]+
  (DXwHat11[:,:,:,2:nz+1]+DYwHat21[:,:,:,2:nz+1]-
  DZuuHat31[:,:,:,2:nz+1]).*wCG[:,:,:,2:nz+1])
FuHat[:,:,:,:,2]= Vort3.*v1CG+
  (DZwwHat32).*wCCG +
  0.5*((DXwHat12[:,:,:,1:nz]  +DYwHat22[:,:,:,1:nz] -
    DZuuHat32[:,:,:,1:nz]).*wCG[:,:,:,1:nz]+
  (DXwHat12[:,:,:,2:nz+1]+DYwHat22[:,:,:,2:nz+1] -
    DZuuHat32[:,:,:,2:nz+1]).*wCG[:,:,:,2:nz+1])

DZwwHat31=(DZwwHat31).*v1CG
DZwwHat32=(DZwwHat32).*v2CG
FuHat[:,:,:,1:nz-1,3]=
  (-DXwHat11[:,:,:,2:nz]-DYwHat21[:,:,:,2:nz]+DZuuHat31[:,:,:,2:nz]).*
    (0.5*(v1CG[:,:,:,1:nz-1]+v1CG[:,:,:,2:nz])) -0.5*
    (DZwwHat31[:,:,:,1:nz-1]+DZwwHat31[:,:,:,2:nz]) +
    (-DXwHat12[:,:,:,2:nz]-DYwHat22[:,:,:,2:nz]+DZuuHat32[:,:,:,2:nz]).*
    (0.5*(v2CG[:,:,:,1:nz-1]+v2CG[:,:,:,2:nz])) -0.5*
    (DZwwHat32[:,:,:,1:nz-1]+DZwwHat32[:,:,:,2:nz])

return FuHat
end