text:
export geam

"""
    geam(alpha::Number, A::CuSparseMatrix, beta::Number, B::CuSparseMatrix, index::SparseChar)

Performs `C = alpha * A + beta * B`. `A` and `B` are sparse matrix defined in CSR storage format.
"""
geam(alpha::Number, A::CuSparseMatrixCSR, beta::Number, B::CuSparseMatrixCSR, index::SparseChar)

for (bname,gname,elty) in ((:cusparseScsrgeam2_bufferSizeExt, :cusparseScsrgeam2, :Float32),
                           (:cusparseDcsrgeam2_bufferSizeExt, :cusparseDcsrgeam2, :Float64),
                           (:cusparseCcsrgeam2_bufferSizeExt, :cusparseCcsrgeam2, :ComplexF32),
                           (:cusparseZcsrgeam2_bufferSizeExt, :cusparseZcsrgeam2, :ComplexF64))
    @eval begin
        function geam(alpha::Number, A::CuSparseMatrixCSR{$elty}, beta::Number, B::CuSparseMatrixCSR{$elty}, index::SparseChar)
            m, n = size(A)
            @assert (m, n) == size(B)
            descrA = CuMatrixDescriptor(CUSPARSE_MATRIX_TYPE_GENERAL, CUSPARSE_FILL_MODE_LOWER, CUSPARSE_DIAG_TYPE_NON_UNIT, index)
            descrB = CuMatrixDescriptor(CUSPARSE_MATRIX_TYPE_GENERAL, CUSPARSE_FILL_MODE_LOWER, CUSPARSE_DIAG_TYPE_NON_UNIT, index)
            descrC = CuMatrixDescriptor(CUSPARSE_MATRIX_TYPE_GENERAL, CUSPARSE_FILL_MODE_LOWER, CUSPARSE_DIAG_TYPE_NON_UNIT, index)

            cusparseSetPointerMode(handle(), CUSPARSE_POINTER_MODE_HOST)
            rowPtrC = CuArray{Int32,1}(undef, m+1)

            function bufferSize()
                out = Ref{Csize_t}(1)
                $bname(handle(), m, n,
                    alpha, descrA, nnz(A), nonzeros(A), A.rowPtr, A.colVal,
                    beta, descrB, nnz(B), nonzeros(B), B.rowPtr, B.colVal,
                    descrC, CuArray{$elty,1}(undef, 0), rowPtrC, CuArray{Int32,1}(undef, 0),
                    out)
                return out[]
            end

            C = with_workspace(bufferSize) do buffer
                function get_nnzC(buffer)
                    nnzTotalDevHostPtr = Ref{Cint}(1)
                    cusparseXcsrgeam2Nnz(handle(), m, n,
                        descrA, nnz(A), A.rowPtr, A.colVal,
                        descrB, nnz(B), B.rowPtr, B.colVal,
                        descrC, rowPtrC, nnzTotalDevHostPtr,
                        buffer)
                    return nnzTotalDevHostPtr[]
                end

                nnzC = get_nnzC(buffer)
                colValC = CuArray{Int32,1}(undef, Int(nnzC))
                nzValC = CuArray{$elty,1}(undef, Int(nnzC))

                $gname(handle(), m, n,
                    alpha, descrA, nnz(A), nonzeros(A), A.rowPtr, A.colVal,
                    beta, descrB, nnz(B), nonzeros(B), B.rowPtr, B.colVal,
                    descrC, nzValC, rowPtrC, colValC,
                    buffer)
                return CuSparseMatrixCSR(rowPtrC, colValC, nzValC, (m, n))
            end
            C
        end
    end
end


text:
using FastAI
import FastAI: createhandle, showblock!

function createhandle(backend::ShowMakie; kwargs...)
    fig = Figure(; kwargs..., backend.kwargs...)
    grid = fig[1, 1] = GridLayout()
    return grid
end

function showblock(backend::ShowMakie, block, obs)
    # Calculate resolution based on number of blocks and size for one block
    # A margin of 10% is added for titles
    width = _nblocks(block) * backend.size[2]
    height = round(Int, 1.1 * backend.size[1])

    grid = createhandle(backend, resolution = (width, height))
    fig = grid.parent.parent
    showblock!(grid, backend, block, obs)
    return fig
end

_nblocks(t::Tuple) = sum(_nblocks, t)
_nblocks(b::AbstractBlock) = 1
_nblocks((_, block)::Pair) = _nblocks(block)

function showblock!(grid, backend::ShowMakie, (title, block)::Pair, obs)
    showblock!(grid[1, 1], backend, block, obs)
end

function showblock!(grid, backend::ShowMakie, blocks::Tuple, obss::Tuple)
    header = [block isa Pair ? first(block) : "" for block in blocks]
    blocks = Tuple(block isa Pair ? last(block) : block for block in blocks)

    rowsize!(grid, 1, Makie.Fixed(backend.size[2]))

    # Show blocks in a row
    col = 1
    for (i, (block, obs)) in enumerate(zip(blocks, obss))
        w = _nblocks(block)
        subgrid = grid[1, col:col+w-1] = GridLayout(tellheight = false)
        showblock!(subgrid, backend, block, obs)
        for j = col:col+w-1
            colsize!(grid, j, Makie.Fixed(backend.size[1]))
        end
        col += w
    end

    # Add titles to named blocks
    Makie.Label(grid[0, 1], "", tellwidth = false, textsize = 25)
    for (i, title) in enumerate(header)
        Makie.Label(grid[1, i], title, tellwidth = false, textsize = 25)
    end

end

function showblocks(backend::ShowMakie, block, obss)
    width = _nblocks(block) * backend.size[2]
    height = round(Int, length(obss) * 1.1 * backend.size[1])
    res = (width * 1.2, height * 1.1)

    grid = createhandle(backend, resolution = res)
    fig = grid.parent.parent

    showblocks!(grid, backend, block, obss)
    return fig
end

function showblocks!(grid, backend::ShowMakie, blocks::Tuple, obss::AbstractVector)
    header = [block isa Pair ? first(block) : "" for block in blocks]
    blocks = Tuple(block isa Pair ? last(block) : block for block in blocks)
    n = _nblocks(blocks)

    # Show each sample in one row
    for (i, obs) in enumerate(obss)
        subgrid = grid[i, 1:n] = GridLayout(tellheight = false)
        rowsize!(grid, i, Makie.Fixed(backend.size[2]))
        showblock!(subgrid, backend, blocks, obs)
    end

    for i = 1:n
        colsize!(grid, i, Makie.Fixed(backend.size[1]))
    end

    # Add titles to named blocks
    Makie.Label(grid[0, 1], "", textsize = 25)
    col = 1
    for (i, (title, block)) in enumerate(zip(header, blocks))
        w = _nblocks(block)
        Makie.Label(grid[1, col:col+w-1], title, tellwidth = false, textsize = 25)
        col += w
    end
end
showblocks!(grid, backend::ShowMakie, block, obss::AbstractVector) =
    showblocks!(grid, backend, (block,), map(obs -> (obs,), obss))

## Block definitions

function showblock!(grid, ::ShowMakie, block::Label, obs)
    ax = cleanaxis(grid[1, 1])
    text!(ax, string(obs), space = :data)
end

function showblock!(grid, ::ShowMakie, block::LabelMulti, obs)
    ax = cleanaxis(grid[1, 1])
    text!(ax, join(string.(obs), "\n"), space = :data)
end

function showblock!(
    grid,
    ::ShowMakie,
    block::Union{<:OneHotTensor{0},<:OneHotTensorMulti{0}},
    obs,
)
    if !(sum(obs) ≈ 1)
        obs = softmax(obs)
    end
    ax = Axis(grid[1, 1], yticks = (1:length(block.classes), string.(block.classes)))
    barplot!(ax, obs, direction = :x)
    hidespines!(ax)
end

function default_showbackend()
    if ismissing(Makie.current_backend[])
        return ShowText()
    else
        return ShowMakie()
    end
end

## Helpers

"""
    cleanaxis(f)

Create a `Makie.Axis` with no interactivity, decorations and aspect distortion.
"""
function cleanaxis(f; kwargs...)
    ax = Makie.Axis(f; kwargs...)
    ax.aspect = Makie.DataAspect()
    ax.xzoomlock = true
    ax.yzoomlock = true
    ax.xrectzoom = false
    ax.yrectzoom = false
    ax.panbutton = nothing
    ax.xpanlock = true
    ax.ypanlock = true
    ax.bottomspinevisible = false
    ax.leftspinevisible = false
    ax.rightspinevisible = false
    ax.topspinevisible = false
    Makie.hidedecorations!(ax)

    return ax
end

ext: jl


text:
module CircuitComponentRounding

import ESeriesRounding: E3, E6, E12, E24, E48, E96, E192
import ESeriesRounding: determine_E

export E3, E6, E12, E24, E48, E96, E192
export determine_E

end


text:
struct UncertainAgeValueDataset <: UncertainIndexValueDataset
    indices::Vector{AbstractUncertainValue}
    values::Vector{AbstractUncertainValue}
end

export UncertainAgeValueDataset


text:
"""
Holds data for hierarchical logistic model test example.
"""
struct BinomLLData{Ty, T}
    y::Ty
    X::T
    Z::T
end
"""
Outer constructor for subsetting by random effects group
"""
function BinomLLData(bld::BinomLLData, grp::Integer)
    idx = Z[:, grp] .== 1
    BinomLLData(bld.y[idx], bld.X[idx, :], bld.Z[idx, grp:grp])
end
"""
Functor that holds parameter values and data for logistic hierarchical
regression model.
"""
struct BinomLL{Tβ, Tu, Tσgrp, Tbld<:BinomLLData} <: Function
    β::Tβ
    u::Tu
    σgrp::Tσgrp
    bld::Tbld
end
"""
Convenient outer constructor to make swapping random effects vectors easier.
"""
function BinomLL(BLL::BinomLL, u)
    BinomLL(BLL.β, u, BLL.σgrp, BLL.bld)
end
"""
Hierarchical logistic regression negative log-likelihood.
"""
function (BLL::BinomLL)()
    lo = BLL.bld.X * BLL.β + BLL.bld.Z * BLL.u
    p = logistic.(lo)
    datalik = -sum(logpdf.(Bernoulli.(p), BLL.bld.y))
    randlik = -sum(logpdf.(Normal(0.0, BLL.σgrp), BLL.u))
    datalik + randlik
end
function (BLL::BinomLL)(u)
    BinomLL(BLL, u)()
end


text:
using OrdinaryDiffEq
using Trixi

###############################################################################
# semidiscretization of the compressible Euler equations

equations = CompressibleEulerEquations2D(1.4)

initial_condition = initial_condition_isentropic_vortex
solver = DGSEM(polydeg=3, surface_flux=flux_lax_friedrichs)

coordinates_min = (-10.0, -10.0)
coordinates_max = ( 10.0,  10.0)
cells_per_dimension = (16, 16)
mesh = CurvedMesh(cells_per_dimension, coordinates_min, coordinates_max)

semi = SemidiscretizationHyperbolic(mesh, equations, initial_condition, solver)

###############################################################################
# ODE solvers, callbacks etc.

tspan = (0.0, 20.0)
ode = semidiscretize(semi, tspan)

summary_callback = SummaryCallback()

analysis_interval = 100
analysis_callback = AnalysisCallback(semi, interval=analysis_interval, save_analysis=true,
                                     extra_analysis_errors=(:conservation_error,),
                                     extra_analysis_integrals=(entropy, energy_total,
                                                               energy_kinetic, energy_internal))

alive_callback = AliveCallback(analysis_interval=analysis_interval)

callbacks = CallbackSet(summary_callback, analysis_callback, alive_callback)

###############################################################################
# run the simulation

sol = solve(ode, BS3(),
            save_everystep=false, callback=callbacks, maxiters=1e5);
summary_callback() # print the timer summary


text:
module Domain

using Dates

export Item

struct Item
  created::DateTime
  name::String
  nbr::Int64
  #constructors
  Item(name, nbr) = new(now(), name, nbr)
end # Item

end


text:
function hasintersection( it::IntervalCollection{S}, i::Interval{T} ) where {S,T}
   !haskey( it.trees, i.seqname ) && return false
   try
      f = first(IntervalTrees.intersect(it.trees[i.seqname], i))
      return true
   catch
      return false
   end
end

function loadbed( filename::String; expandfirst::Int=0, expandlast::Int=0 )
   r  = BED.Reader(open( filename ))
   ic = IntervalCollection{String}()
   for b in r
      push!( ic, Interval{String}(seqname(b),
                                  leftposition(b) - expandfirst,
                                  rightposition(b) + expandlast,
                                  strand(b), String(b.data[b.name])) )
   end
   ic
end

# Eventually replace with Automa.jl parser...
function loadbedgraph( filename::String; regions::IntervalCollection{T}=IntervalCollection{Void}(), regionbool=false, expand=0 ) where T
   fh = open( filename , "r")
   if isgzipped( filename )
      fh = fh |> x->ZlibInflateInputStream(x, reset_on_end=true)
   end

   retcol = IntervalCollection{Float64}()

   for i in readlines(fh)
      ismatch( r"^#", i ) && continue
      spl = split(i, '\t')
      first = parse(Int, spl[2]) + 1
      last  = parse(Int, spl[3])
      if regionbool && !hasintersection( regions, Interval(spl[1], first-expand, last+expand) )
         continue
      else
         score = parse(Float64, spl[4])
         push!( retcol, Interval{Float64}(spl[1], first, last, '?', score) )
      end
   end

   retcol
end


text:
module Col

using Compat
using DataFrames
import Iterators
import Base: ==

immutable GroupedColumn
    columns::Nullable{Vector}
end

function Base.hash(colgroup::GroupedColumn, h::UInt64)
    return hash(colgroup.columns, h)
end

function ==(a::GroupedColumn, b::GroupedColumn)
    return (isnull(a.columns) && isnull(b.columns)) ||
        (!isnull(a.columns) && !isnull(b.columns) && get(a.columns) == get(b.columns))
end

function Base.show(io::IO, gc::GroupedColumn)
    print(io, "Column")
end

function index()
    return GroupedColumn(Nullable{Vector}())
end

function index{T <: (@compat Union{Int, Symbol})}(xs::T...)
    return GroupedColumn(Nullable(collect(T, xs)))
end

immutable GroupedColumnValue
    columns::Nullable{Vector}
end

function Base.show(io::IO, gc::GroupedColumnValue)
    print(io, "Column Value")
end

function value()
    return GroupedColumnValue(Nullable{Vector}())
end

function value{T <: (@compat Union{Int, Symbol})}(xs::T...)
    return GroupedColumnValue(Nullable(collect(T, xs)))
end

end # module Col

# Handle aesthetics aliases and warn about unrecognized aesthetics.
#
# Returns:
#   A new mapping with aliases evaluated and unrecognized aesthetics removed.
#
function cleanmapping(mapping::Dict)
    cleaned = Dict{Symbol, Any}()
    for (key, val) in mapping
        # skip the "order" pesudo-aesthetic, used to order layers
        if key == :order
            continue
        end

        if haskey(aesthetic_aliases, key)
            key = aesthetic_aliases[key]
        elseif !in(key, fieldnames(Aesthetics))
            warn("$(string(key)) is not a recognized aesthetic. Ignoring.")
            continue
        end

        if val == Col.value || val == Col.index
            val = val()
        end

        cleaned[key] = val
    end
    cleaned
end

# Type to contain fields produced by melting data (and to dispatch on)
immutable MeltedData
    data
    melted_data
    indicators::Array
    colmap::Dict
end

function meltdata(U::AbstractDataFrame, colgroups_::Vector{Col.GroupedColumn})
    um, un = size(U)

    colgroups = Set(colgroups_)

    # Figure out the size of the new melted matrix
    allcolumns = Set{Symbol}(names(U))

    vm = um
    grouped_columns = Set{Symbol}()
    for colgroup in colgroups
        if isnull(colgroup.columns) # null => group all columns
            vm *= un
            grouped_columns = copy(allcolumns)
        else
            for j in get(colgroup.columns)
                if !isa(j, Symbol)
                    error("DataFrame columns can only be grouped by (Symbol) names")
                end
                push!(grouped_columns, j)
            end
            vm *= length(get(colgroup.columns))
        end
    end

    ungrouped_columns = setdiff(allcolumns, grouped_columns)
    vn = length(colgroups) + length(ungrouped_columns)

    V = AbstractArray[]
    vnames = Symbol[]
    colmap = Dict{Any, Int}()

    # allocate vectors for grouped columns
    for (j, colgroup) in enumerate(colgroups)
        cols = isnull(colgroup.columns) ? allcolumns : get(colgroup.columns)

        # figure the grouped common column type
        firstcol = U[first(cols)]
        eltyp = eltype(firstcol)
        vectyp = isa(firstcol, Vector) ? Vector : DataVector
        for col in cols
            eltyp = promote_type(eltyp, typeof(U[col]))
            if !isa(U[col], Vector)
                vectyp = DataVector
            end
        end

        push!(V, eltyp == Vector ? Array(eltyp, vm) : DataArray(eltyp, vm))
        name = gensym()
        push!(vnames, name)
        colmap[colgroup] = j
    end

    # allocate vectors for ungrouped columns
    for (j, col) in enumerate(ungrouped_columns)
        push!(V, similar(U[col], vm))
        colmap[col] = j + length(colgroups)
        push!(vnames, col)
    end

    # Indicator columns for each colgroup
    indicators = Array(Symbol, (vm, length(colgroups)))

    colidxs = [isnull(colgroup.columns) ? collect(allcolumns) : get(colgroup.columns)
               for colgroup in colgroups]

    vi = 1
    for ui in 1:um
        for colidx in Iterators.product(colidxs...)
            # copy grouped columns
            for (vj, uj) in enumerate(colidx)
                V[vj][vi] = U[ui, uj]
                indicators[vi, vj] = uj
            end

            # copy ungrouped columns
            for (vj, uj) in enumerate(ungrouped_columns)
                V[vj + length(colgroups)][vi] = U[ui, uj]
            end

            vi += 1
        end
    end

    df = DataFrame(; collect(zip(vnames, V))...)
    return MeltedData(U, df, indicators, colmap)
end

# TODO: The is too elaborate. All matrix melts should const of all column, and
# we should reserve this elaborate melting logic for data frames.
function meltdata(U::AbstractMatrix, colgroups_::Vector{Col.GroupedColumn})
    um, un = size(U)

    colgroups = Set(colgroups_)

    # Figure out the size of the new melted matrix
    allcolumns = IntSet()
    for i in 1:un
        push!(allcolumns, i)
    end

    vm = um
    grouped_columns = IntSet()
    for colgroup in colgroups
        if isnull(colgroup.columns)
            vm *= un
            grouped_columns = copy(allcolumns)
        else
            for j in get(colgroup.columns)
                push!(grouped_columns, j)
            end
            vm *= length(get(colgroup.columns))
        end
    end

    ungrouped_columns = setdiff(allcolumns, grouped_columns)
    vn = length(colgroups) + length(ungrouped_columns)
    V = similar(U, (vm, vn))

    # Indicator columns for each colgroup
    indicators = Array(Int, (vm, length(colgroups)))

    colidxs = [isnull(colgroup.columns) ? collect(allcolumns) : get(colgroup.columns)
               for colgroup in colgroups]

    vi = 1
    for ui in 1:um
        for colidx in Iterators.product(colidxs...)
            # copy grouped columns
            for (vj, uj) in enumerate(colidx)
                V[vi, vj] = U[ui, uj]
                indicators[vi, vj] = uj
            end

            # copy ungrouped columns
            for (vj, uj) in enumerate(ungrouped_columns)
                V[vi, vj + length(colgroups)] = U[ui, uj]
            end

            vi += 1
        end
    end

    # Map grouped and individual columns in U to columns in V
    colmap = Dict{Any, Int}()
    for (vj, colgroup) in enumerate(colgroups)
        colmap[colgroup] = vj
    end
    for (vj, uj) in enumerate(ungrouped_columns)
        colmap[uj] = vj + length(colgroups)
    end

    return MeltedData(U, V, indicators, colmap)
end

# Evaluate one mapping.
evalmapping(source, arg::AbstractArray) = arg
evalmapping(source, arg::Function) = arg
evalmapping(source, arg::Distribution) = arg

evalmapping(source::AbstractDataFrame, arg::Symbol) = source[arg]
evalmapping(source::AbstractDataFrame, arg::AbstractString) = evalmapping(source, Symbol(arg))
evalmapping(source::AbstractDataFrame, arg::Integer) = source[arg]
evalmapping(source::AbstractDataFrame, arg::Expr) = with(source, arg)

function evalmapping(source::MeltedData, arg::Integer)
    return source.melted_data[:,source.colmap[arg]]
end

function evalmapping(source::MeltedData, arg::Col.GroupedColumn)
    return source.indicators[:,source.colmap[arg]]
end

function evalmapping(source::MeltedData, arg::Col.GroupedColumnValue)
    return source.melted_data[:,source.colmap[Col.GroupedColumn(arg.columns)]]
end

function evalmapping(source::MeltedData, arg::Colon)
    return source.melted_data
end

# Evalute aesthetic mappings producting a Data instance.
function evalmapping!(mapping::Dict, data_source, data::Data)
    # Are we doing implicit reshaping?
    colgroups = Col.GroupedColumn[]
    for (k, v) in mapping
        if isa(v, Col.GroupedColumn)
            push!(colgroups, v)
        elseif isa(v, Col.GroupedColumnValue)
            push!(colgroups, Col.GroupedColumn(v.columns))
        end
    end

    if !isempty(colgroups)
        data_source = meltdata(data_source, colgroups)
    end

    for (k, v) in mapping
        setfield!(data, k, evalmapping(data_source, v))
        data.titles[k] = isa(v, AbstractString) || isa(v, Symbol) ?  string(v) : string(k)
    end

    return data_source
end


text:
using nng
using Test

function test_push_pull()
    push_s = nng.listen("ipc:///tmp/pipeline.ipc", nng.PUSH0)
    pull_s = nng.dial("ipc:///tmp/pipeline.ipc", nng.PULL0)

    nng.send(push_s, "Something")
    result = nng.recv(pull_s)

    nng.close(push_s)
    nng.close(pull_s)
    return result
end

function test_pub_sub()
    pub_s = nng.listen("ipc:///tmp/pub.ipc", nng.PUB0)
    sub_s = nng.dial("ipc:///tmp/pub.ipc", nng.SUB0, "") # Subscribe to all topics
    sleep(1)

    nng.send(pub_s, "Something")
    result = nng.recv(sub_s)

    nng.close(pub_s)
    nng.close(sub_s)
    return result
end

function test_pair()
    pair_1 = nng.listen("ipc:///tmp/pair.ipc", nng.PAIR0)
    pair_2 = nng.dial("ipc:///tmp/pair.ipc", nng.PAIR0)
    sleep(1)

    nng.send(pair_1, "One way")
    result = nng.recv(pair_2)

    nng.close(pair_1)
    nng.close(pair_2)

    return result
end

function test_req_rep()
    rep = nng.listen("ipc:///tmp/reqrep.ipc", nng.REP0)
    req = nng.dial("ipc:///tmp/reqrep.ipc", nng.REQ0)

    nng.send(req, "One way")
    message = nng.recv(rep)

    nng.send(rep, "The other way")
    message = nng.recv(req)

    nng.close(rep)
    nng.close(req)

    return message
end

function test_surveyor_respondent()
    surv = nng.listen("ipc:///tmp/survey.ipc", nng.SURVEYOR0)
    resp = nng.dial("ipc:///tmp/survey.ipc", nng.RESPONDENT0)
    sleep(1)

    nng.send(surv, "One way")
    message = nng.recv(resp)

    nng.send(resp, "The other way")
    message = nng.recv(surv)

    nng.close(surv)
    nng.close(resp)

    return message
end

function test_bus()
    bus1 = nng.listen("ipc:///tmp/bus1.ipc", nng.BUS0)
    bus2 = nng.listen("ipc:///tmp/bus2.ipc", nng.BUS0)
    sleep(1)

    bus1 = nng.dial(bus1, "ipc:///tmp/bus2.ipc")
    bus2 = nng.dial(bus2, "ipc:///tmp/bus1.ipc")
    sleep(1)

    nng.send(bus1, "One way")
    message = nng.recv(bus2)

    nng.close(bus1)
    nng.close(bus2)

    return message
end

@testset "nng.jl" begin
    @test test_push_pull() == "Something"
    @test test_pub_sub() == "Something"
    @test test_pair() == "One way"
    @test test_req_rep() == "The other way"
    @test test_surveyor_respondent() == "The other way"
    @test test_bus() == "One way"
end


text:
grad_mut(b::Buffer) = fill!(similar(b.data, Any), nothing)
grad_mut(b::Buffer{T}) where T<:Number = fill!(similar(b.data, float(T)), 0)

@nograd Buffer

@adjoint function getindex(b::Buffer, i...)
  b[i...], function (Δ)
    grad = grad_mut(__context__, b)
    grad[i...] = accum(grad[i...], Δ)
    return
  end
end

@adjoint! function setindex!(b::Buffer, v, i...)
  setindex!(b, v, i...), function (_)
    grad = grad_mut(__context__, b)
    v̄ = grad[i...]
    zero = eltype(grad) <: Number ? 0 : nothing
    if i isa NTuple{N,Integer} where N
      grad[i...] = zero
    else
      grad[i...] .= zero
    end
    (nothing, v̄, map(_->nothing, i)...)
  end
end

@adjoint! function copyto!(b::Buffer, xs)
  copyto!(b, xs), function (_)
    grad = grad_mut(__context__, b)
    x̄s = copy(grad)
    grad .= eltype(grad) <: Number ? 0 : nothing
    return (nothing, x̄s)
  end
end

@adjoint! function push!(b::Buffer, x)
  push!(b, x), function (y)
    grad = grad_mut(__context__, b)
    return (nothing, pop!(grad))
  end
end

_pullback(cx::AContext, ::typeof(Broadcast.materialize!), b::Buffer, x::AbstractArray) =
  _pullback(cx, copyto!, b, x)

@adjoint function copy(b::Buffer)
  res = copy(b)

  function copy_sensitivity(b̄)
    grad_mut(__context__, b)[:] .= vec(b̄)
    return
  end

  function copy_sensitivity(b̄::Union{Tuple,AbstractVector{T}}) where {T<:AbstractArray}
    grad_mut(__context__, b)[:] .= b̄
    return
  end

  return res, copy_sensitivity
end


text:
using Base.Threads

# v = [[1, 2, 3], [1, 2, 3], [1, 2, 3],  [1, 2, 5], [1, 3, 4], [1, 4, 5], [2, 3, 4], [2, 3, 4], [2, 3, 5], [3, 4, 5]]
# v = [sample(1:20, 5) for x = 1:10_000];

function combine_dicts!(d, d2)
    all_keys = [collect(keys(d)); collect(keys(d2))]
    uniq_keys = unique(all_keys)

    for k in uniq_keys
        d[k] = get(d, k, 0) + get(d2, k, 0)
    end
end

# Given an integer and a number of chunks, this function
# returns a vector of ranges that partition 1:n into
# nchnks total chunks.
function gen_ranges(n, nchnks)
    step_size = round(Int64, (ceil(n/nchnks)))
    first_idx = 1
    last_idx = 1 + step_size
    range_vec = Array{Range, 1}(nchnks)
    for i = 1:nchnks
        range_vec[i] = range(first_idx, step_size)
        first_idx += step_size
        last_chk = first_idx + step_size
        step_size = last_chk > n ? (n - first_idx + 1) : step_size
    end
    return range_vec
end

function gen_ranges(first_idx, n, nchnks)
    num = n - first_idx + 1
    step_size = round(Int64, (ceil(num/nchnks)))

    last_idx = 1 + step_size
    range_vec = Array{Range, 1}(nchnks)
    for i = 1:nchnks
        range_vec[i] = range(first_idx, step_size)
        first_idx += step_size
        last_chk = first_idx + step_size
        step_size = last_chk > n ? (n - first_idx + 1) : step_size
    end
    return range_vec
end

# Given a vector of vectors, this function returns a
# Dict with counts representing the frequency of each
# individual value.
function count_values(u)
    out = Dict()
    for i = 1:length(u)
        for v in u[i]
            out[v] = get(out, v, 0) + 1
        end
    end
    out
end

function pcount_values(v)
    nt = nthreads()
    dict_vec = Vector{Dict}(nt)
    range_vec = gen_ranges(length(v), nt)

    @threads for j = 1:nt
        dict_vec[j] = count_values(v[range_vec[j]])
    end
    cnt_dict = Dict()

    for j = 2:nt
        combine_dicts!(dict_vec[1], dict_vec[j])
    end
    return dict_vec[1]
end

# @time count_values(v)
# @time pcount_values(v)
#
# d1 = count_values(v[1:50_000])
# d2 = count_values(v[50_001:100_000])

function init_nested_vectors(typ, num, num_inner_elem = 0)
    inner = Array{typ, 1}(num_inner_elem)
    outer = Array{Array, 1}(0)

    for i = 1:num
        push!(outer, copy(inner))
    end
    outer
end

# init_nested_vectors(Array, 4)

# This version of the function seems to work, but because of
# an issue with the type of `i` being inferred as `Box.Core`,
# the performance is no better than the serial implentation. This
# should be fixed once the Github issue #15276 is resolved.

function par_apriori_gen{M}(x::Array{Array{M, 1}, 1}, num_threads)
    n = length(x)
    m = length(x[1]) - 1
    C_arr = init_nested_vectors(Array, num_threads)

    for i = 1:n
        sort!(x[i])
        range_vec = gen_ranges(i+1, n, num_threads)

        @threads for t = 1:num_threads
            for j in range_vec[t]
                sort!(x[j])
                keep_candidate = true

                # length k candidate itemsets are created by merging pairs of
                # length k - 1 itemsets if their first k - 2 elements identical
                for l in 1:m

                    # see if all k - 1 elements are identical
                    if x[i][l] != x[j][l] || x[i][m+1] == x[j][m+1]
                        keep_candidate = false
                        break
                    end
                end
                if keep_candidate
                    c::Array{M, 1} = vcat(x[i], x[j][end])
                    push!(C_arr[t], sort!(c))
                end
            end
        end
    end
    C::Array{Array{M, 1}, 1} = reduce(vcat, C_arr)
    return C              # vector of candidate itemsets: C_{k}
end

function par_apriori_gen2{M}(x::Array{Array{M, 1}, 1}, num_threads)
    n = length(x)

    m = length(x[1]) - 1
    C_arr = Array{Array{M, 1}, 1}(0)

    for i = 1:n
        sort!(x[i])
        range_vec = gen_ranges(i+1, n, num_threads)

        @parallel for t in 1:num_threads
            for j in range_vec[t]

                sort!(x[j])
                keep_candidate = true

                # length k candidate itemsets are created by merging pairs of
                # length k - 1 itemsets if their first k - 2 elements identical
                for l in 1:m

                    # see if all k - 1 elements are identical
                    if x[i][l] != x[j][l] || x[i][m+1] == x[j][m+1]
                        keep_candidate = false
                        break
                    end
                end
                if keep_candidate
                    c::Array{M, 1} = vcat(x[i], x[j][end])
                    push!(C_arr[t], sort!(c))
                end
            end
        end
    end
    return C_arr              # vector of candidate itemsets: C_{k}
end

# v = [sample(1:20, 5, replace = false) for x = 1:10_000];

function _par_apriori_gen{M}(x::Array{Array{M, 1}, 1}, num_threads, i, range_vec)
    n = length(x)
    m = length(x[1]) - 1
    C_arr = init_nested_vectors(Array, num_threads)
    sort!(x[i])

    @threads for t = 1:num_threads
        for j in range_vec[t]
            sort!(x[j])
            keep_candidate = true

            # length k candidate itemsets are created by merging pairs of
            # length k - 1 itemsets if their first k - 2 elements identical
            for l in 1:m

                # see if all k - 1 elements are identical
                if x[i][l] != x[j][l] || x[i][m+1] == x[j][m+1]
                    keep_candidate = false
                    break
                end
            end
            if keep_candidate
                c::Array{M, 1} = vcat(x[i], x[j][end])
                push!(C_arr[t], sort!(c))
            end
        end
    end

    C::Array{Array{M, 1}, 1} = reduce(vcat, C_arr)
    return C              # vector of candidate itemsets: C_{k}
end

function par_apriori_gen3{M}(x::Array{Array{M, 1}, 1}, num_threads, n)
    C_arr = Array{Array{Array{M, 1}, 1}}(n)
    for i = 1:n
        range_vec = gen_ranges(i+1, n, num_threads)
        C_arr[i] = _par_apriori_gen(x, num_threads, i, range_vec)
    end
    res = reduce(vcat, C_arr)
end


text:
# License for this file: MIT (expat)
# Copyright 2017-2018, DLR Institute of System Dynamics and Control

"""
    module ModiaMath.NonlinearEquations

This module contains functions to solve nonlinear algebraic equations:

- [`solveOneNonlinearEquation`](@ref):
  Function that computes the solution of one non-linear algebraic equation `y=f(u)`
  in one unknown `u` in a reliable and efficient way using Brents algorithm.

- [`KINSOL`](@ref):
  Module containing functions to solve a system of nonlinear systems of equations
  with Sundials KINSOL. The module is designed so that the same system is
  solved several times with KINSOL as needed by Modia simulations
  (auxiliary memory is only allocated once and not for every call).
  KINSOL is used in ModiaMath to solve nonlinear algebraic equations during
  initialization and at events of a simulation.

# Main developer
[Martin Otter](https://rmc.dlr.de/sr/de/staff/martin.otter/),
[DLR - Institute of System Dynamics and Control](https://www.dlr.de/sr/en)
"""
module NonlinearEquations

export KINSOL, NonlinearEquationsInfo, solveNonlinearEquations!
export solveOneNonlinearEquation

# using/imports
import ModiaMath

# include code
include("solveOneNonlinearEquation.jl")
include("sundials_kinsol.jl")

end


text:
abstract Player

type Human <: Player
end

#=type Network <: Player=#
#=end=#

function get_valid_move(h::Human, b::Board, w::Ptr{Void})
    m = Char(TermWin.wgetch(w))
    while !is_valid(m,b)
        m = Char(TermWin.wgetch(w))
    end
    return m
end

function key_to_command(k::Char)
    if k == 'h'
        "w"
    elseif k == 'j'
        "s"
    elseif k == 'k'
        "n"
    elseif k == 'l'
        "e"
    elseif k == 'y'
        "nw"
    elseif k == 'u'
        "ne"
    elseif k == 'b'
        "sw"
    elseif k == 'n'
        "se"
    elseif k == 'w'
        "wait"
    else
        string(k)
    end
end

function execute_command(com::AbstractString,b::Board)
    if com == "t"
        teleport_sprite!(b)
    elseif com == " "
        skip_sprite(b)
    elseif com == "wait"
        enter_wait_mode!(b)
    else
        move_sprite!(b,com)
    end
end

max_stars_repo_name: RatanRSur/robots-from-scratch

max_issues_repo_name: RatanRSur/robots-from-scratch

max_forks_repo_name: RatanRSur/robots-from-scratch


text:
function filt_zeros!(zoomer::AbstractArray{T}) where {T<:AbstractFloat}
    for i in eachindex(zoomer)
        if isapprox(zoomer[i], zero(T), atol = eps(T))
            zoomer[i] = one(T)
        end
    end
end

struct LinearScaler{N,V<:AbstractVector} <: Scaler{N}
    offset::V
    zoomer::V
end

function StandardScaler(
    mean::AbstractVector{T},
    std::AbstractVector{T},
) where {T<:AbstractFloat}
    if length(mean) != length(std)
        throw(DimensionMismatch("`mean` and `std` must have the same length"))
    end
    μ = copy(mean)
    σ = copy(std)
    filt_zeros!(σ)
    LinearScaler{length(μ),typeof(μ)}(μ, σ)
end

function StandardScaler(x::AbstractMatrix{T}) where {T<:AbstractFloat}
    n, m = size(x)
    if m < 2
        throw(ArgumentError("`x` must have at least 2 samples (cols)"))
    end
    μ = similar(x, n)
    mean!(μ, x)
    σ = reshape(std(x, dims = 2, mean = μ), :)
    filt_zeros!(σ)
    LinearScaler{length(μ),typeof(μ)}(μ, σ)
end

function MinMaxScaler(
    min::AbstractVector{T},
    max::AbstractVector{T},
) where {T<:AbstractFloat}
    if length(min) != length(max)
        throw(DimensionMismatch("`min` and `max` must have the same length"))
    end
    dlt = max - min
    filt_zeros!(dlt)
    LinearScaler{length(dlt),typeof(dlt)}(copy(min), dlt)
end

function MinMaxScaler(x::AbstractMatrix{T}) where {T<:AbstractFloat}
    n, m = size(x)
    if m < 2
        throw(ArgumentError("`x` must have at least 2 samples (cols)"))
    end
    min = similar(x, n)
    minimum!(min, x)
    max = similar(min)
    maximum!(max, x)
    @. max -= min
    filt_zeros!(max)
    LinearScaler{length(max),typeof(max)}(min, max)
end

transform!(s::LinearScaler, x::AbstractVecOrMat) =
    @. x = (x - s.offset) / s.zoomer

transform(s::LinearScaler, x::AbstractVecOrMat) = transform!(s, copy(x))

inverse!(s::LinearScaler, x::AbstractVecOrMat) =
    @. x = (x * s.zoomer) + s.offset

inverse(s::LinearScaler, x::AbstractVecOrMat) = inverse!(s, copy(x))

Base.:*(s::LinearScaler, x::AbstractVecOrMat) = transform(s, x)

Base.:\(s::LinearScaler, x::AbstractVecOrMat) = inverse(s, x)

inverse_variance!(s::LinearScaler, v::AbstractVecOrMat) = @. v *= s.zoomer^2
transform_variance!(s::LinearScaler, v::AbstractVecOrMat) = @. v /= s.zoomer^2

Base.eltype(s::LinearScaler) = eltype(s.offset)


text:
using ITensors,
      Test

@testset "MPO Basics" begin
  N = 6
  sites = SiteSet(N,2)
  @test length(MPO()) == 0
  O = MPO(sites)
  @test length(O) == N

  str = split(sprint(show, O), '\n')
  @test str[1] == "MPO"
  @test length(str) == length(O) + 2

  O[1] = ITensor(sites[1], prime(sites[1]))
  @test hasindex(O[1],sites[1])
  @test hasindex(O[1],prime(sites[1]))

  P = copy(O)
  @test hasindex(P[1],sites[1])
  @test hasindex(P[1],prime(sites[1]))

  @testset "orthogonalize" begin
    phi = randomMPS(sites)
    K = randomMPO(sites)
    orthogonalize!(phi, 1)
    orthogonalize!(K, 1)
    orig_inner = inner(phi, K, phi)
    orthogonalize!(phi, div(N, 2))
    orthogonalize!(K, div(N, 2))
    @test inner(phi, K, phi) ≈ orig_inner
  end

  @testset "inner" begin
    phi = randomMPS(sites)
    K = randomMPO(sites)
    @test maxLinkDim(K) == 1
    psi = randomMPS(sites)
    phidag = dag(phi)
    prime!(phidag)
    phiKpsi = phidag[1]*K[1]*psi[1]
    for j = 2:N
      phiKpsi *= phidag[j]*K[j]*psi[j]
    end
    @test phiKpsi[] ≈ inner(phi,K,psi)

    badsites = SiteSet(N+1,2)
    badpsi = randomMPS(badsites)
    @test_throws DimensionMismatch inner(phi,K,badpsi)
  end

  @testset "applyMPO" begin
    phi = randomMPS(sites)
    K = randomMPO(sites)
    @test maxLinkDim(K) == 1
    psi = randomMPS(sites)
    psi_out = applyMPO(K, psi,maxdim=1)
    @test inner(phi,psi_out) ≈ inner(phi,K,psi)
    @test_throws ArgumentError applyMPO(K, psi, method="fakemethod")

    badsites = SiteSet(N+1,2)
    badpsi = randomMPS(badsites)
    @test_throws DimensionMismatch applyMPO(K,badpsi)
  end
  @testset "add" begin
    shsites = spinHalfSites(N)
    K = randomMPO(shsites)
    L = randomMPO(shsites)
    M = sum(K, L)
    @test length(M) == N
    psi = randomMPS(shsites)
    k_psi = applyMPO(K, psi, maxdim=1)
    l_psi = applyMPO(L, psi, maxdim=1)
    @test inner(psi, sum(k_psi, l_psi)) ≈ inner(psi, M, psi) atol=5e-3
  end

  @testset "nmultMPO" begin
    psi = randomMPS(sites)
    K = randomMPO(sites)
    L = randomMPO(sites)
    @test maxLinkDim(K) == 1
    @test maxLinkDim(L) == 1
    KL = nmultMPO(K, L, maxdim=1)
    psi_kl_out = applyMPO(K, applyMPO(L, psi, maxdim=1), maxdim=1)
    @test inner(psi,KL,psi) ≈ inner(psi, psi_kl_out) atol=5e-3

    badsites = SiteSet(N+1,2)
    badL = randomMPO(badsites)
    @test_throws DimensionMismatch nmultMPO(K,badL)
  end

  sites = spinHalfSites(N)
  O = MPO(sites,"Sz")
  @test length(O) == N # just make sure this works

  @test_throws ArgumentError randomMPO(sites, 2)
  @test_throws ErrorException linkindex(MPO(N, fill(ITensor(), N), 0, N + 1), 1)
end


text:
using Nemo, GibbsTypePriors, RCall, JLD

grid_k = 1:100

n, β, σ = 100, 1.2, 0.6

exact_Pkn = GibbsTypePriors.Pkn_NGG_arb.(grid_k, n, β, σ)
approx_Pkn = GibbsTypePriors.Pkn_NGG_approx.(grid_k, n, β, σ, n, exact_Pkn[n])

R"library(tidyverse)"
R"p2 = tibble(k = $grid_k,
         exact_Pkn = $(Float64.(exact_Pkn)),
         approx_Pkn = $(Float64.(approx_Pkn))) %>%
         #mutate(abs_rel_error = abs(approx_Pkn-exact_Pkn)/exact_Pkn) %>%
         mutate(Error = approx_Pkn-exact_Pkn) %>%
         #ggplot(aes(x = k, y = abs_rel_error)) +
         ggplot(aes(x = k, y = Error)) +
         theme_bw() +
         geom_point() +
         geom_line()"

R"p1 = tibble(k = $grid_k,
        exact_Pkn = $(Float64.(exact_Pkn)),
        approx_Pkn = $(Float64.(approx_Pkn))) %>%
        gather(type, value, -k) %>%
        mutate(type = gsub('_Pkn', '', type)) %>%
        ggplot(aes(x = k, y = value, colour = type, group = type)) +
        theme_bw() +
        geom_point() +
        geom_line() +
        theme(legend.position = 'left') +
        scale_colour_discrete(name='') +
        ylab('Pkn')"

R"p = gridExtra::grid.arrange(p1, p2, nrow = 1)"

R"pdf('test/graphical_tests/figures_graphical_tests/Pkn_NGG_approx_quality.pdf', height = 5, width = 7)
        plot(p)
dev.off()"

max_stars_repo_name: konkam/GibbsTypePriors

max_issues_repo_name: konkam/GibbsTypePriors

max_forks_repo_name: konkam/GibbsTypePriors


text:
include("9.1 Integrand simplification rules.jl")
include("9.2 Derivative integration rules.jl")
include("9.3 Piecewise linear functions.jl")
include("9.4 Miscellaneous integration rules.jl")

max_stars_repo_path: src/rules/9 Miscellaneous/.jl

max_issues_repo_path: src/rules/9 Miscellaneous/.jl

max_forks_repo_path: src/rules/9 Miscellaneous/.jl


text:
"""
    The Hull is defined by :
    * Φ: angle of rotation of the tank around the longitudinal axis. If the right side of the tank is lower than the left side,
    than this angle is positif.
    * Ψ: angle defined by the longitudinal axes of the tank and the horizontal plane
    If the nose of the tank is higher than the back, this angle is positif
    * Χ: angle defined by the the projection of the longitudinal axis of the tank into an horizontal plane
    and the X axis
"""
mutable struct Hull
    Φ::Union{Float64,Nothing}
    ψ::Union{Float64,Nothing}
    χ::Union{Float64,Nothing}
end
"""
    The turret is defined by its azimuth angle
"""
mutable struct Turret
     ξ::Union{Float64,Nothing}#azimuth
end
"""
    The Canon is defined by:
    * θ elevation angle (degrees)
    * lw the length of the canon
    * u₀ the muzzle velocity magnitude
"""
mutable struct Canon
     θ::Union{Float64,Nothing} #elevation
     lw::Float64
     u₀::Float64
     twist::Float64
end
"""
    The sight is defined by:
    * θ LOS elevation angle
    * ξ LOS azimuth angle
"""
mutable struct Sight
    θ::Float64
    ξ::Float64
end
"""
    The tank is defined by: a hull, a turret, a canon, a sight and its position (altitude)
"""
mutable struct Tank
   hull::Hull
   turret::Turret
   canon::Canon
   sight::Sight
   rWY::Float64
   latitude::Float64
   position::Union{Vector{Float64},Nothing}
   velocity::Union{Vector{Float64},Nothing}
   function Tank(hull::Hull,turret::Turret,canon::Canon,sight::Sight,rWY::Float64,latitude::Float64)
    new(hull,turret,canon,sight,rWY,latitude,nothing,nothing)
end
end


text:
using LoopVectorization
using LinearAlgebra
using TensorOperations

function RMP2(Alg::RMP2Algorithm)
    aoints = IntegralHelper{Float64}()
    rhf = RHF(aoints)

    if typeof(aoints.eri_type) === JKFIT || Options.get("precision") == "single"
        aoints = IntegralHelper(eri_type=RIFIT())
    end
    moints = IntegralHelper(orbitals=rhf.orbitals)
    RMP2(moints, aoints, Alg)
end

function RMP2(aoints::IntegralHelper{Float64,E,AtomicOrbitals}, Alg::RMP2Algorithm) where E <: AbstractERI
    rhf = RHF(aoints)
    moints = IntegralHelper(orbitals=rhf.orbitals)
    RMP2(moints, aoints, Alg)
end

function RMP2(O::AbstractRestrictedOrbitals, Alg::RMP2Algorithm)
    moints = IntegralHelper(orbitals=O)
    aoints = IntegralHelper(eri_type=moints.eri_type)
    RMP2(moints, aoints, Alg)
end

function RMP2(rhf::RHF, Alg::RMP2Algorithm)
    moints = IntegralHelper(orbitals=rhf.orbitals)
    aoints = IntegralHelper(eri_type=moints.eri_type)
    RMP2(moints, aoints, Alg)
end

function RMP2(M::Molecule, Alg::RMP2Algorithm)
    aoints = IntegralHelper{Float64}(molecule = M, eri_type=SparseERI)
    rhf = RHF(aoints)
    moints = IntegralHelper(molecule = M, orbitals=rhf.orbitals)
    RMP2(moints, aoints, Alg)
end

function RMP2(moints::IntegralHelper{T,Chonky,<:AbstractRestrictedOrbitals}, aoints::IntegralHelper{T,<:AbstractERI,AtomicOrbitals}, Alg::RMP2Algorithm) where T<:AbstractFloat
    mp_header()
    Fermi.Integrals.compute!(moints, aoints, "OVOV")
    Fermi.Integrals.compute!(moints, aoints, "F")
    RMP2(moints, Alg)
end

function RMP2(moints::IntegralHelper{T,E,<:AbstractRestrictedOrbitals}, aoints::IntegralHelper{T,E,AtomicOrbitals}, Alg::RMP2Algorithm) where {T<:AbstractFloat, E<:AbstractDFERI}
    mp_header()
    Fermi.Integrals.compute!(moints, aoints, "BOV")
    Fermi.Integrals.compute!(moints, aoints, "F")
    RMP2(moints, Alg)
end

function RMP2(ints::IntegralHelper{<:AbstractFloat,<:AbstractERI,<:AbstractRestrictedOrbitals}, Alg::RMP2Algorithm)

    # Check frozen core and inactive virtual
    ndocc = ints.molecule.Nα
    nvir = size(ints.orbitals.C,1) - ndocc
    core = Options.get("drop_occ")
    inac = Options.get("drop_vir")

    if core ≥ ndocc
        throw(FermiException("invalid number of frozen orbitals ($core) for $ndocc doubly occupied orbitals"))
    end
    if inac ≥ nvir
        throw(FermiException("invalid number of inactive virtual orbitals ($inac) for $nvir total virtual orbitals"))
    end

    output("  Starting MP2 computation")
    output(" Number of frozen orbitals:             {:d}" , core)
    output(" Number of inactive orbitals:           {:d}" , inac)
    output(" Number of correlated electron pairs:   {:d}", ndocc-core)
    output(" Number of correlated virtual orbitals: {:d}", nvir-inac)
    output(" ⇒ Total number of MP2 amplitudes:      {:d}", (ndocc-core)^2*(nvir-inac)^2)

    output(repeat("-",80))

    # Compute MP2 energy
    Emp2 = RMP2_energy(ints, Alg)
    Eref = ints.orbitals.sd_energy

    output("   @Final RMP2 Correlation Energy {:>20.12f} Eₕ", Emp2)
    output("   Reference Energy               {:>20.12f} Eₕ", Eref)
    output("   @Final RMP2 Total Energy       {:>20.12f} Eₕ", Emp2+Eref)
    output(repeat("-",80))

    RMP2(Emp2, Emp2+Eref)
end

function RMP2_energy(ints::IntegralHelper{T, <:AbstractDFERI, RHFOrbitals}, Alg::RMP2Algorithm) where T<:AbstractFloat
    Bvo = permutedims(ints["BOV"], (1,3,2))
    ϵo = ints["Fii"]
    ϵv = ints["Faa"]

    output(" Computing DF-MP2 Energy!")
    v_size = length(ϵv)
    o_size = length(ϵo)

    # Pre-allocating arrays for threads
    BABs = [zeros(T, v_size, v_size) for _ = 1:Threads.nthreads()]

    # Disable BLAS threading
    nt  = BLAS.get_num_threads()
    BLAS.set_num_threads(1)

    # Vector containing the energy contribution computed by each thread
    ΔMP2s = zeros(T,Threads.nthreads())
    TWO = T(2.0)
    ONE = one(T)
    t = @elapsed begin
        @sync for i in 1:o_size
            Threads.@spawn begin
            id = Threads.threadid()

            Bab = BABs[id]
            @views Bi = Bvo[:,:,i]

            for j in i:o_size

                @views Bj = Bvo[:,:,j]
                mul!(Bab, transpose(Bi), Bj)

                eij = ϵo[i] + ϵo[j]
                E = zero(T)
                @inbounds for a = eachindex(ϵv)
                    eija = eij - ϵv[a]
                    for b = eachindex(ϵv)
                        D = eija - ϵv[b]
                        E += Bab[a,b] * (TWO * Bab[a,b] - Bab[b,a]) / D
                    end
                end
                fac = i !== j ? TWO : ONE
                ΔMP2s[id] += fac * E
            end
        end # spawn
        end # sync
        Emp2 = sum(ΔMP2s)
    end # time
    output("Done in {:5.5f} seconds.", t)

    BLAS.set_num_threads(nt)
    return Emp2
end

function RMP2_energy(ints::IntegralHelper{T, Chonky, RHFOrbitals}, Alg::RMP2Algorithm) where T<:AbstractFloat
    output(" Computing MP2 Energy... ", ending="")
    ovov = ints["OVOV"]
    ϵo = ints["Fii"]
    ϵv = ints["Faa"]

    ΔMP2s = zeros(T,Threads.nthreads())
    TWO = T(2.0)
    t = @elapsed begin
        @sync for b in eachindex(ϵv)
            Threads.@spawn begin
            id = Threads.threadid()
            for a in eachindex(ϵv)
                @inbounds ϵ_ab = ϵv[a] + ϵv[b]
                for j in eachindex(ϵo)
                    @inbounds ϵ_abj = ϵo[j] - ϵ_ab
                    for i in eachindex(ϵo)
                        @fastmath @inbounds ΔMP2s[id] += ovov[i,a,j,b]*(TWO*ovov[i,a,j,b] - ovov[i,b,j,a]) / (ϵo[i]+ϵ_abj)
                    end
                end
            end
            end
        end
    end
    output("Done in {:5.5f} s\n", t)
    return sum(ΔMP2s)
end

function RMP2_energy(ints::IntegralHelper{T, Chonky, <:AbstractRestrictedOrbitals}, Alg::RMP2Algorithm) where T<:AbstractFloat
    output(" Computing Iterative MP2 Energy")

    foo = ints["Fij"]
    fvv = ints["Fab"]
    ϵo = ints["Fii"]
    ϵv = ints["Faa"]
    Vovov = ints["OVOV"]
    newT2 = permutedims(Vovov, (1,3,2,4))
    invD = [1.0/(ϵo[i]+ϵo[j]-ϵv[a]-ϵv[b]) for i=eachindex(ϵo), j=eachindex(ϵo), a=eachindex(ϵv), b=eachindex(ϵv)]
    T2 = similar(newT2)

    # Iteration parameters
    ite = 1
    Emp2 = 0.0
    oldE = 0.0
    dE = 1.0
    rms = 1.0
    main_time = 0.0

    e_conv = Options.get("cc_e_conv")
    max_iter = Options.get("cc_max_iter")
    max_rms = Options.get("cc_max_rms")

    output("{:10s}    {: 15s}    {: 12s}    {:12s}    {:10s}","Iteration","CC Energy","ΔE","Max RMS","time (s)")
    while (abs(dE) > e_conv || rms > max_rms)
        if ite > max_iter
            output("\n⚠  MP2 Equations did not converge in {:1.0d} iterations.", max_iter)
            break
        end
        t = @elapsed begin

            T2 .= newT2
            oldE = Emp2
            TWO = T(2)
            @tensor begin
                newT2[i,j,a,b]  = fvv[c,a]*T2[i,j,c,b]
                newT2[i,j,a,b] -= foo[i,k]*T2[k,j,a,b]
            end
            newT2 .+= permutedims(newT2, (2,1,4,3))
            newT2 .+= permutedims(Vovov, (1,3,2,4))
            newT2 .*= invD

            @tensor begin
                #Energy
                B[l,c,k,d] := TWO*newT2[k,l,c,d]
                B[l,c,k,d] -= newT2[l,k,c,d]
                Emp2 = B[l,c,k,d]*Vovov[k,c,l,d]
            end

            # Compute residues
            rms = sqrt(sum((newT2 .- T2).^2)/length(T2))
        end
        dE = Emp2 - oldE
        main_time += t
        output("    {:<5.0d}    {:<15.10f}    {:>+12.10f}    {:<12.10f}    {:<10.5f}", ite, Emp2, dE, rms, t)
        ite += 1
    end
    output("\nMain MP2 iterations done in {:5.5f} s", main_time)
    output("Average time per iteration {:5.5f}", main_time/(ite-1))

    # Converged?
    conv = false
    if abs(dE) < e_conv && rms < max_rms
        output("\n 🍾 Equations Converged!")
        conv = true
    end

    return Emp2
end

max_stars_repo_name: FermiQC/Fermi

max_issues_repo_name: FermiQC/Fermi

max_forks_repo_name: FermiQC/Fermi


text:
using Clustering;
using Distances;

function getLandmarks(X, p, method=:Random)
    if(method == :Random)
        numberOfPoints = size(X,2);
        landmarks = X[:,randperm(numberOfPoints)[1:p]];
        return landmarks;
    end

    if(method == :Kmeans)
        kmeansResult = kmeans(X,p)
        return kmeansResult.centers;
    end

    throw(ArgumentError("method can only be :Kmeans or :Random"));
end

function gaussianKernel(distance, bandwidth)
    exp(-distance / (2*bandwidth^2));
end

function getLinearCoding(X, landmarks, bandwidth, r)
    distances = pairwise(Distances.Euclidean(), landmarks, X);
    similarities = map(x -> gaussianKernel(x, bandwidth), distances);
    ZHat = spzeros(size(similarities,1), size(similarities,2));

    for i in 1:(size(similarities,2))
        topLandMarksIndices = selectperm(similarities[:,i], 1:r, rev=true);
        topLandMarksCoefficients = similarities[topLandMarksIndices, i];
        topLandMarksCoefficients = topLandMarksCoefficients / sum(topLandMarksCoefficients);
        ZHat[topLandMarksIndices,i] = topLandMarksCoefficients;
    end
    return diagm(sum(ZHat,2)[:])^(-1/2) * ZHat;
end

function LSCClustering(X, nrOfClusters, nrOfLandmarks, method, nonZeroLandmarkWeights, bandwidth)
    landmarks = getLandmarks(X, nrOfLandmarks, method)
    ZHat = getLinearCoding(X, landmarks, bandwidth, nonZeroLandmarkWeights)
    svdResult = svd(transpose(ZHat))
    clusteringResult = kmeans(transpose(svdResult[1][:,1:nrOfClusters]),nrOfClusters);
    return clusteringResult
end


text:
function _check_dimension(v::AbstractVector, s)
    length(v) != MOI.dimension(s) && throw(DimensionMismatch("Mismatch between value and set"))
    return nothing
end

function distance_to_set(::DefaultDistance, v::AbstractVector{T}, s::MOI.Reals) where {T <: Real}
    _check_dimension(v, s)
    return zero(T)
end

function distance_to_set(::NormedEpigraphDistance{p}, v::AbstractVector{<:Real}, s::MOI.Zeros) where {p}
    _check_dimension(v, s)
    return LinearAlgebra.norm(v, p)
end

function distance_to_set(::DefaultDistance, v::AbstractVector{<:Real}, s::MOI.Zeros)
    return distance_to_set(NormedEpigraphDistance{2}(), v, s)
end

function distance_to_set(::NormedEpigraphDistance{p}, v::AbstractVector{<:Real}, s::MOI.Nonnegatives) where {p}
    _check_dimension(v, s)
    return LinearAlgebra.norm((ifelse(vi < 0, -vi, zero(vi)) for vi in v), p)
end

function distance_to_set(::DefaultDistance, v::AbstractVector{<:Real}, s::MOI.Nonnegatives)
    return distance_to_set(NormedEpigraphDistance{2}(), v, s)
end

function distance_to_set(::NormedEpigraphDistance{p}, v::AbstractVector{<:Real}, s::MOI.Nonpositives) where {p}
    _check_dimension(v, s)
    return LinearAlgebra.norm((ifelse(vi > 0, vi, zero(vi)) for vi in v), p)
end

function distance_to_set(::DefaultDistance, v::AbstractVector{<:Real}, s::MOI.Nonpositives)
    _check_dimension(v, s)
    return distance_to_set(NormedEpigraphDistance{2}(), v, s)
end

distance_to_set(::EpigraphViolationDistance, v::Real, s::MOI.LessThan) = max(v - s.upper, zero(v))
distance_to_set(::EpigraphViolationDistance, v::Real, s::MOI.GreaterThan) = max(s.lower - v, zero(v))
distance_to_set(::EpigraphViolationDistance, v::Real, s::MOI.EqualTo) = abs(v - s.value)

distance_to_set(::DefaultDistance, v, s::Union{MOI.LessThan, MOI.GreaterThan, MOI.EqualTo}) = distance_to_set(EpigraphViolationDistance(), v, s)

distance_to_set(::DefaultDistance, v::T, s::MOI.Interval) where {T <: Real} = max(s.lower - v, v - s.upper, zero(T))

function distance_to_set(::EpigraphViolationDistance, v::AbstractVector{<:Real}, s::MOI.NormInfinityCone)
    _check_dimension(v, s)
    t = v[1]
    xs = @view(v[2:end])
    result = maximum(abs, xs) - t
    return max(result, zero(result))
end

function distance_to_set(::EpigraphViolationDistance, v::AbstractVector{<:Real}, s::MOI.NormOneCone)
    _check_dimension(v, s)
    t = v[1]
    xs = @view(v[2:end])
    result = sum(abs, xs) - t
    return max(result, zero(result))
end

function distance_to_set(::EpigraphViolationDistance, v::AbstractVector{<:Real}, s::MOI.SecondOrderCone)
    _check_dimension(v, s)
    t = v[1]
    xs = @view(v[2:end])
    result = LinearAlgebra.norm2(xs) - t
    return max(result, zero(result))
end

"""
    distance_to_set(::NormedEpigraphDistance{p}, v::AbstractVector{<:Real}, s::MOI.SecondOrderCone)

Composition of the projection with the p-norm.
"""
function distance_to_set(::NormedEpigraphDistance{p}, v::AbstractVector{<:Real}, s::MOI.SecondOrderCone) where {p}
    _check_dimension(v, s)
    t = v[1]
    xs = @view(v[2:end])
    nx = LinearAlgebra.norm(xs, p)
    if nx <= t
        return zero(eltype(nx))
    end
    if nx <= -t
        return sqrt(t^2 + nx^2)
    end
    return (nx - t) / √2
end

function distance_to_set(::DefaultDistance, v, s::MOI.SecondOrderCone)
    return distance_to_set(NormedEpigraphDistance{2}(), v, s)
end

function distance_to_set(::DefaultDistance, v, s::Union{MOI.NormInfinityCone, MOI.NormOneCone})
    return distance_to_set(EpigraphViolationDistance(), v, s)
end

function distance_to_set(::NormedEpigraphDistance{p}, v::AbstractVector{<:Real}, s::MOI.RotatedSecondOrderCone) where {p}
    _check_dimension(v, s)
    t = v[1]
    u = v[2]
    xs = @view(v[3:end])
    return LinearAlgebra.norm(
        (max(-t, zero(t)), max(-u, zero(u)), max(LinearAlgebra.dot(xs,xs) - 2 * t * u)),
        p,
    )
end

function distance_to_set(::DefaultDistance, v::AbstractVector{<:Real}, s::MOI.RotatedSecondOrderCone)
    return distance_to_set(NormedEpigraphDistance{2}(), v, s)
end

function distance_to_set(::DefaultDistance, v::AbstractVector{<:Real}, s::MOI.GeometricMeanCone)
    _check_dimension(v, s)
    t = v[1]
    xs = @view(v[2:end])
    n = MOI.dimension(s) - 1
    xresult = LinearAlgebra.norm2(
        max.(-xs, zero(eltype(xs)))
    )
    # early returning if there exists x[i] < 0 to avoid complex sqrt
    if xresult > 0
        return xresult
    end
    result = t - prod(xs)^(inv(n))
    return max(result, zero(result))
end

function distance_to_set(::NormedEpigraphDistance{p}, v::AbstractVector{<:Real}, s::MOI.ExponentialCone) where {p}
    _check_dimension(v, s)
    x = v[1]
    y = v[2]
    z = v[3]

    if x <= 0 && isapprox(y, 0, atol=1e-10) && z >= 0
        return zero(eltype(v))
    end

    result = y * exp(x/y) - z
    return LinearAlgebra.norm(
        (max(-y, zero(result)), max(result, zero(result))),
        p,
    )
end

function distance_to_set(::DefaultDistance, v::AbstractVector{<:Real}, s::MOI.ExponentialCone)
    return distance_to_set(NormedEpigraphDistance{2}(), v, s)
end

function distance_to_set(::NormedEpigraphDistance{p}, vs::AbstractVector{<:Real}, s::MOI.DualExponentialCone) where {p}
    _check_dimension(vs, s)
    u = vs[1]
    v = vs[2]
    w = vs[3]

    if isapprox(u, 0, atol=1e-10) && v >= 0 && w >= 0
        return zero(eltype(vs))
    end

    result = -u*exp(v/u) - ℯ * w
    return LinearAlgebra.norm(
        (max(u, zero(result)), max(result, zero(result))),
        p,
    )
end

function distance_to_set(::DefaultDistance, vs::AbstractVector{<:Real}, s::MOI.DualExponentialCone)
    return distance_to_set(NormedEpigraphDistance{2}(), vs, s)
end

function distance_to_set(::DefaultDistance, v::AbstractVector{<:Real}, s::MOI.PowerCone)
    _check_dimension(v, s)
    x = v[1]
    y = v[2]
    z = v[3]
    e = s.exponent
    # early return to avoid complex exponent results
    if x < 0 || y < 0
       return LinearAlgebra.norm2(
            (max(-x, zero(x)), max(-y, zero(x)))
        )
    end
    result = abs(z) - x^e * y^(1-e)
    return max(result, zero(result))
end

function distance_to_set(::NormedEpigraphDistance{p}, vs::AbstractVector{<:Real}, s::MOI.DualPowerCone) where {p}
    _check_dimension(vs, s)
    u = vs[1]
    v = vs[2]
    w = vs[3]
    e = s.exponent
    ce = 1-e
    result = abs(w) - (u/e)^e * (v/ce)^ce
    return LinearAlgebra.norm(
        (max(-u, zero(result)), max(-v, zero(result)), max(result, zero(result))),
        p,
    )
end

function distance_to_set(::DefaultDistance, vs::AbstractVector{<:Real}, s::MOI.DualPowerCone)
    return distance_to_set(NormedEpigraphDistance{2}(), vs, s)
end

function distance_to_set(::NormedEpigraphDistance{p}, v::AbstractVector{<:Real}, set::MOI.RelativeEntropyCone) where {p}
    _check_dimension(v, s)
    n = (dimension(set)-1) ÷ 2
    u = v[1]
    v = @view(v[2:(n+1)])
    w = @view(v[(n+2):end])
    s = sum(w[i] * log(w[i]/v[i]) for i in eachindex(w))
    result = s - u
    return LinearAlgebra.norm(
        push!(
            max.(@view(v[2:end]), zero(result)),
            max(result, zero(result)),
        ),
        p,
    )
end

function distance_to_set(::DefaultDistance, v::AbstractVector{<:Real}, set::MOI.RelativeEntropyCone)
    return distance_to_set(NormedEpigraphDistance{2}(), v, set)
end

function distance_to_set(::EpigraphViolationDistance, v::AbstractVector{<:Real}, s::MOI.NormSpectralCone)
    _check_dimension(v, s)
    t = v[1]
    m = reshape(@view(v[2:end]), (s.row_dim, s.column_dim))
    s1 = LinearAlgebra.svd(m).S[1]
    result = s1 - t
    return max(result, zero(result))
end

function distance_to_set(::EpigraphViolationDistance, v::AbstractVector{<:Real}, s::MOI.NormNuclearCone)
    _check_dimension(v, s)
    t = v[1]
    m = reshape(@view(v[2:end]), (s.row_dim, s.column_dim))
    s1 = sum(LinearAlgebra.svd(m).S)
    result = s1 - t
    return max(result, zero(result))
end

function distance_to_set(::DefaultDistance, v, s::Union{MOI.NormSpectralCone, MOI.NormNuclearCone})
    return distance_to_set(EpigraphViolationDistance(), v, s)
end

function distance_to_set(::DefaultDistance, v::T, ::MOI.ZeroOne) where {T <: Real}
    return min(abs(v - zero(T)), abs(v - one(T)))
end

function distance_to_set(::DefaultDistance, v::Real, ::MOI.Integer)
    return min(abs(v - floor(v)), abs(v - ceil(v)))
end

# return the element-wise distance to zero, with the greatest element to 0
function distance_to_set(::DefaultDistance, v::AbstractVector{T}, ::MOI.SOS1) where {T <: Real}
    _check_dimension(v, s)
    _, i = findmax(abs.(v))
    return LinearAlgebra.norm2([v[j] for j in eachindex(v) if j != i])
end

# takes in input [z, f(x)]
function distance_to_set(::DefaultDistance, v::AbstractVector{T}, s::MOI.IndicatorSet{A}) where {A, T <: Real}
    _check_dimension(v, s)
    z = v[1]
    # inactive constraint
    if A === MOI.ACTIVATE_ON_ONE && isapprox(z, 0) || A === MOI.ACTIVATE_ON_ZERO && isapprox(z, 1)
        return zeros(T, 2)
    end
    return LinearAlgebra.norm2(
        (distance_to_set(d, z, MOI.ZeroOne()), distance_to_set(v[2], s.set))
    )
end

function distance_to_set(::NormedEpigraphDistance{p}, v::AbstractVector{T}, s::MOI.PositiveSemidefiniteConeTriangle) where {p, T <: Real}
    X = unvec_symm(v, s.side_dimension)
    λ, U = LinearAlgebra.eigen(X)
    Tp = eltype(λ)
    λm = -min.(λ, zero(Tp))
    vdist = vec_symm(U * LinearAlgebra.Diagonal(λm) * U')
    return LinearAlgebra.norm(vdist, p)
end

function distance_to_set(::DefaultDistance, v::AbstractVector{T}, s::MOI.PositiveSemidefiniteConeTriangle) where {T <: Real}
    return distance_to_set(NormedEpigraphDistance{2}(), v, s)
end


text:
module TestClustering

using MLJBase
using Test
using Random:seed!
import LinearAlgebra: norm
import Distances: evaluate
using CategoricalArrays

# load code to be tested:
import MLJModels
import Clustering # MLJModels.Clustering_ now available for loading
using MLJModels.Clustering_

seed!(132442)

X, y = @load_crabs

####
#### KMEANS
####

barekm = KMeans()

fitresult, cache, report = fit(barekm, 1, X)

R = matrix(transform(barekm, fitresult, X))

X_array = matrix(X)

# distance from first point to second center
@test R[1, 2] ≈ norm(view(X_array, 1, :) .- view(fitresult[1], :, 2))^2
@test R[10, 3] ≈ norm(view(X_array, 10, :) .- view(fitresult[1], :, 3))^2

p = predict(barekm, fitresult, X)

@test argmin(R[1, :]) == p[1]
@test argmin(R[10, :]) == p[10]

infos = info_dict(barekm)

@test infos[:package_name] == "Clustering"
@test infos[:is_pure_julia]
@test infos[:package_license] == "MIT"

@test infos[:input_scitype] == Table(Continuous)
@test infos[:output_scitype] == Table(Continuous)

infos[:docstring]

####
#### KMEDOIDS
####

barekm = KMedoids()

fitresult, cache, report = fit(barekm, 1, X)

R = matrix(transform(barekm, fitresult, X))

@test R[1, 2] ≈ evaluate(barekm.metric, view(X_array, 1, :), view(fitresult[1], :, 2))
@test R[10, 3] ≈ evaluate(barekm.metric, view(X_array, 10, :), view(fitresult[1], :, 3))

p = predict(barekm, fitresult, X)

@test all(report.assignments .== p)

# km = machine(barekm, X)
# fit!(km)

infos = info_dict(barekm)

@test infos[:input_scitype] == Table(Continuous)
@test infos[:output_scitype] == Table(Continuous)

end # module
true


text:
@testset "DIIS" begin
    dm = Fermi.DIIS.DIISManager{Float32, Float32}()
    @test length(dm) == 0
    x = rand(10)
    y = rand(10)
    push!(dm, x, y)
    @test length(dm) == 1
    @test_throws BoundsError Fermi.DIIS.extrapolate(dm)
end


text:
n_per_category(s, h, t, y, n_in_t, n_occurence_history) = (n_in_t[t] ∩ n_occurence_history[s, h, y])

function γ_s!(γ_s, γ, n_all)
    K, D, size_memory, T, rain_cat = size(γ_s)
    for tup in Iterators.product(1:D, 1:size_memory, 1:T, 1:rain_cat)
        for k = 1:K
            γ_s[k, tup...] = sum(γ[n, k] for n in n_all[tup...]; init = 0)
        end
    end
end

function s_ξ!(s_ξ, ξ, n_in_t)
    T, K = size(s_ξ, 1), size(s_ξ, 2)
    for t = 1:T
        for (k, l) in Iterators.product(1:K, 1:K)
            s_ξ[t, k, l] = sum(ξ[n, k, l] for n in n_in_t[t])
        end
    end
    # * We add ξ[N, k, l] but it should be zeros
end

function model_for_B_d(γ_s::AbstractMatrix, d::Int; silence = true)

    T, rain_cat = size(γ_s)
    model = Model(Ipopt.Optimizer)
    set_optimizer_attribute(model, "max_cpu_time", 60.0)
    set_optimizer_attribute(model, "max_iter", 100)

    silence && set_silent(model)
    f = 2π / T
    cos_nj = [cos(f * j * t) for t = 1:T, j = 1:d]
    sin_nj = [sin(f * j * t) for t = 1:T, j = 1:d]
    trig = [[1; interleave2(cos_nj[t, :], sin_nj[t, :])] for t = 1:T]

    @variable(model, θ_jump[j = 1:(2d+1)])
    # Polynomial P
    @NLexpression(model, Pn[t = 1:T], sum(trig[t][j] * θ_jump[j] for j = 1:length(trig[t])))

    @NLparameter(model, π_s[t = 1:T, y = 1:rain_cat] == γ_s[t, y])

    @NLexpression(model, mle,
        -sum(π_s[t, 1] * log1p(exp(-Pn[t])) for t = 1:T) - sum(π_s[t, 2] * log1p(exp(+Pn[t])) for t = 1:T)
    ) # 1 is where it did not rain # 2 where it rained
    @NLobjective(
        model, Max,
        mle
    )
    # I don't know if it is the best but https://discourse.julialang.org/t/jump-updating-nlparameter-of-a-model-in-a-loop/35081/3
    model[:π_s] = π_s
    return model
end

function update_B_d!(B::AbstractArray{T,4} where {T}, θ_Y::AbstractArray{N,4} where {N}, γ::AbstractMatrix, γ_s::AbstractArray, observations, n_all, model_B::Model; warm_start = true)
    @argcheck size(γ, 1) == size(observations, 1)
    @argcheck size(γ, 2) == size(B, 1)
    N = size(γ, 1)
    K = size(B, 1)
    T = size(B, 2)
    D = size(B, 3)
    size_memory = size(B, 4)
    ## For periodicHMM only the n observations corresponding to B(t) are used to update B(t)
    ## Update the smoothing parameters in the JuMP model

    γ_s!(γ_s, γ, n_all) # update coefficient in JuMP model

    θ_res = pmap(tup -> fit_mle_one_B_d(θ_Y[tup..., :], model_B, γ_s[tup..., :, :]; warm_start = warm_start), Iterators.product(1:K, 1:D, 1:size_memory))

    for (k, s, h) in Iterators.product(1:K, 1:D, 1:size_memory)
        θ_Y[k, s, h, :] = θ_res[k, s, h]
    end

    p = [1 / (1 + exp(polynomial_trigo(t, θ_Y[k, s, h, :], T = T))) for k = 1:K, t = 1:T, s = 1:D, h = 1:size_memory]
    B[:, :, :, :] = Bernoulli.(p)
end

function fit_mle_one_B_d(θ_Y, model_B, γ_s; warm_start = true)
    T, rain_cat = size(γ_s)
    θ_jump = model_B[:θ_jump]
    warm_start && set_start_value.(θ_jump, θ_Y[:])
    π_s = model_B[:π_s]

    for t = 1:T, y = 1:rain_cat
        set_value(π_s[t, y], γ_s[t, y])
    end
    optimize!(model_B)
    value.(θ_jump)
end

# JuMP model use to increase R(θ,θ^i) for the Q(t) matrix
function model_for_A_d(s_ξ::AbstractArray, d::Int; silence = true)
    T, K = size(s_ξ)
    model = Model(Ipopt.Optimizer)
    set_optimizer_attribute(model, "max_iter", 200)
    silence && set_silent(model)
    f = 2π / T
    cos_nj = [cos(f * j * t) for t = 1:T, j = 1:d]
    sin_nj = [sin(f * j * t) for t = 1:T, j = 1:d]

    trig = [[1; interleave2(cos_nj[t, :], sin_nj[t, :])] for t = 1:T]

    @variable(model, pklj_jump[l = 1:(K-1), j = 1:(2d+1)], start = 0.01)
    # Polynomial P_kl
    @NLexpression(model, Pkl[t = 1:T, l = 1:K-1], sum(trig[t][j] * pklj_jump[l, j] for j = 1:length(trig[t])))

    @NLparameter(model, s_πkl[t = 1:T, l = 1:K-1] == s_ξ[t, l])
    #TODO? is it useful to define the extra parameter for the sum?
    @NLparameter(model, s_πk[t = 1:T] == sum(s_ξ[t, l] for l = 1:K))

    @NLobjective(
        model,
        Max,
        sum(sum(s_πkl[t, l] * Pkl[t, l] for l = 1:K-1) - s_πk[t] * log1p(sum(exp(Pkl[t, l]) for l = 1:K-1)) for t = 1:T)
    )
    # To add NL parameters to the model for later use https://discourse.julialang.org/t/jump-updating-nlparameter-of-a-model-in-a-loop/35081/3
    model[:s_πkl] = s_πkl
    model[:s_πk] = s_πk
    return model
end

function update_A_d!(
    A::AbstractArray{T,3} where {T},
    θ_Q::AbstractArray{T,3} where {T},
    ξ::AbstractArray,
    s_ξ::AbstractArray,
    α::AbstractMatrix,
    β::AbstractMatrix,
    LL::AbstractMatrix,
    n2t::AbstractArray{Int},
    n_in_t,
    model_A::Model;
    warm_start = true
)
    @argcheck size(α, 1) == size(β, 1) == size(LL, 1) == size(ξ, 1)
    @argcheck size(α, 2) ==
              size(β, 2) ==
              size(LL, 2) ==
              size(A, 1) ==
              size(A, 2) ==
              size(ξ, 2) ==
              size(ξ, 3)

    N, K = size(LL)
    T = size(A, 3)
    @inbounds for n in OneTo(N - 1)
        t = n2t[n] # periodic t
        m = SHHMM.vec_maximum(view(LL, n + 1, :)) #TODO Should be imported from HMMBase maxmouchet
        c = 0.0

        for i in OneTo(K), j in OneTo(K)
            ξ[n, i, j] = α[n, i] * A[i, j, t] * exp(LL[n+1, j] - m) * β[n+1, j]
            c += ξ[n, i, j]
        end

        for i in OneTo(K), j in OneTo(K)
            ξ[n, i, j] /= c
        end
    end
    ##
    # ξ are the filtering probablies
    s_ξ!(s_ξ, ξ, n_in_t)
    θ_res = pmap(k -> fit_mle_one_A_d(θ_Q[k, :, :], model_A, s_ξ[:, k, :]; warm_start = warm_start), 1:K)

    for k = 1:K
        θ_Q[k, :, :] = θ_res[k][:, :]
    end

    #TODO change the vectorize form to loop (or something else) to avoid allocation
    [A[k, l, t] = exp(polynomial_trigo(t, θ_Q[k, l, :], T = T)) for k = 1:K, l = 1:K-1, t = 1:T]
    [A[k, K, t] = 1 for k = 1:K, t = 1:T] # last colum is 1/normalization (one could do otherwise)
    normalization_polynomial = [1 + sum(A[k, l, t] for l = 1:K-1) for k = 1:K, t = 1:T]
    [A[k, l, t] /= normalization_polynomial[k, t] for k = 1:K, l = 1:K, t = 1:T]
end

function fit_mle_one_A_d(θ_Q, model, s_ξ; warm_start = true)
    T, K = size(s_ξ)
    pklj_jump = model[:pklj_jump]
    s_πk = model[:s_πk]
    s_πkl = model[:s_πkl]
    ## Update the smoothing parameters in the JuMP model
    for t = 1:T
        set_value(s_πk[t], sum(s_ξ[t, l] for l = 1:K))
        for l = 1:K-1
            set_value(s_πkl[t, l], s_ξ[t, l])
        end
    end
    warm_start && set_start_value.(pklj_jump, θ_Q[:, :])
    # Optimize the updated model
    optimize!(model)
    # Obtained the new parameters
    return value.(pklj_jump)
end

function fit_mle(hmm::HierarchicalPeriodicHMM, observations, n2t::AbstractArray{Int},
    θ_Q::AbstractArray{TQ,3} where {TQ},
    θ_Y::AbstractArray{TY,4} where {TY}; all_iters = false, kwargs...)
    hmm = copy(hmm)
    θ_Q = copy(θ_Q)
    θ_Y = copy(θ_Y)
    history, all_θ_Q, all_θ_Y = fit_mle!(hmm, observations, n2t, θ_Q, θ_Y; kwargs...)
    if all_iters == true
        return hmm, θ_Q, θ_Y, history, all_θ_Q, all_θ_Y
    else
        return hmm, θ_Q, θ_Y, history
    end
end

function fit_mle!(
    hmm::HierarchicalPeriodicHMM,
    observations::AbstractArray,
    n2t::AbstractArray{Int},
    θ_Q::AbstractArray{TQ,3} where {TQ},
    θ_Y::AbstractArray{TY,4} where {TY}
    ;
    display = :none,
    maxiter = 100,
    tol = 1e-3,
    robust = false,
    silence = true,
    warm_start = true,
    past = [0 1 0 1 1 0 1 0 0 0
        1 1 0 1 1 1 1 1 1 1
        1 1 0 1 1 1 0 1 1 1
        1 1 0 1 1 0 0 0 1 0
        1 1 0 1 1 0 0 1 0 1]
)
    @argcheck display in [:none, :iter, :final]
    @argcheck maxiter >= 0

    N, K, T, size_memory, D = size(observations, 1), size(hmm, 1), size(hmm, 3), size(hmm, 4), size(hmm, 2)

    deg_Q = (size(θ_Q, 3) - 1) ÷ 2
    deg_Y = (size(θ_Y, 4) - 1) ÷ 2
    rain_cat = 2
    @argcheck T == size(hmm.B, 2)
    history = EMHistory(false, 0, [])

    all_θ_Q = [copy(θ_Q)]
    all_θ_Y = [copy(θ_Y)]
    # Allocate memory for in-place updates
    c = zeros(N)
    α = zeros(N, K)
    β = zeros(N, K)
    γ = zeros(N, K) # regular smoothing proba
    γ_s = zeros(K, D, size_memory, T, rain_cat) # summed smoothing proba
    ξ = zeros(N, K, K)
    s_ξ = zeros(T, K, K)
    LL = zeros(N, K)

    # assign category for observation depending in the past observations
    memory = Int(log(size_memory) / log(2))
    lag_cat = conditional_to(observations, memory; past = past)

    n_in_t = [findall(n2t .== t) for t = 1:T]
    n_occurence_history = [findall(.&(observations[:, j] .== y, lag_cat[:, j] .== h)) for j = 1:D, h = 1:size_memory, y = 0:1]
    n_all = [n_per_category(tup..., n_in_t, n_occurence_history) for tup in Iterators.product(1:D, 1:size_memory, 1:T, 1:rain_cat)]

    model_A = model_for_A_d(s_ξ[:, 1, :], deg_Q, silence = silence) # JuMP Model for transition matrix
    model_B = model_for_B_d(γ_s[1, 1, 1, :, :], deg_Y, silence = silence) # JuMP Model for Emmission distribution

    loglikelihoods!(LL, hmm, observations, n2t, lag_cat)
    robust && replace!(LL, -Inf => nextfloat(-Inf), Inf => log(prevfloat(Inf)))

    forwardlog!(α, c, hmm.a, hmm.A, LL, n2t)
    backwardlog!(β, c, hmm.a, hmm.A, LL, n2t)
    posteriors!(γ, α, β)

    logtot = sum(c)
    (display == :iter) && println("Iteration 0: logtot = $logtot")

    for it = 1:maxiter
        update_a!(hmm.a, α, β)
        update_A_d!(hmm.A, θ_Q, ξ, s_ξ, α, β, LL, n2t, n_in_t, model_A; warm_start = warm_start)
        update_B_d!(hmm.B, θ_Y, γ, γ_s, observations, n_all, model_B; warm_start = warm_start)
        # Ensure the "connected-ness" of the states,
        # this prevents case where there is no transitions
        # between two extremely likely observations.
        robust && (hmm.A .+= eps())

        @check isprobvec(hmm.a)
        @check istransmats(hmm.A)

        push!(all_θ_Q, copy(θ_Q))
        push!(all_θ_Y, copy(θ_Y))

        # loglikelihoods!(LL, hmm, observations, n2t)
        loglikelihoods!(LL, hmm, observations, n2t, lag_cat)

        robust && replace!(LL, -Inf => nextfloat(-Inf), Inf => log(prevfloat(Inf)))

        forwardlog!(α, c, hmm.a, hmm.A, LL, n2t)
        backwardlog!(β, c, hmm.a, hmm.A, LL, n2t)
        posteriors!(γ, α, β)

        logtotp = sum(c)
        (display == :iter) && println(now(), " Iteration $it: logtot = $logtotp")
        flush(stdout)

        push!(history.logtots, logtotp)
        history.iterations += 1

        if abs(logtotp - logtot) < tol
            (display in [:iter, :final]) &&
                println("EM converged in $it iterations, logtot = $logtotp")
            history.converged = true
            break
        end

        logtot = logtotp
    end

    if !history.converged
        if display in [:iter, :final]
            println("EM has not converged after $(history.iterations) iterations, logtot = $logtot")
        end
    end

    history, all_θ_Q, all_θ_Y
end

# TODO add possibiluity of memory different at each site
function fit_mle!(
    hmm::HierarchicalPeriodicHMM,
    observations::AbstractArray,
    n2t::AbstractArray{Int},
    θ_Q::AbstractArray{TQ,3} where {TQ},
    θ_Y::AbstractArray{TY,4} where {TY},
    size_memories::AbstractVector # Vector of all local memory when there are not indentical
    ;
    display = :none,
    maxiter = 100,
    tol = 1e-3,
    robust = false,
    silence = true,
    warm_start = true,
    past = [0 1 0 1 1 0 1 0 0 0
        1 1 0 1 1 1 1 1 1 1
        1 1 0 1 1 1 0 1 1 1
        1 1 0 1 1 0 0 0 1 0
        1 1 0 1 1 0 0 1 0 1]
)
    @argcheck display in [:none, :iter, :final]
    @argcheck maxiter >= 0

    N, K, T, D = size(observations, 1), size(hmm, 1), size(hmm, 3), size(hmm, 2)
    @argcheck length(size_memories) == D
    max_size_memory = maximum(size_memories)

    deg_Q = (size(θ_Q, 3) - 1) ÷ 2
    deg_Y = (size(θ_Y, 4) - 1) ÷ 2
    rain_cat = 2
    @argcheck T == size(hmm.B, 2)
    history = EMHistory(false, 0, [])

    all_θ_Q = [copy(θ_Q)]
    all_θ_Y = [copy(θ_Y)]
    # Allocate memory for in-place updates
    c = zeros(N)
    α = zeros(N, K)
    β = zeros(N, K)
    γ = zeros(N, K) # regular smoothing proba
    γ_s = zeros(K, D, max_size_memory, T, rain_cat) # summed smoothing proba
    ξ = zeros(N, K, K)
    s_ξ = zeros(T, K, K)
    LL = zeros(N, K)

    # assign category for observation depending in the past observations
    memories = Int.(log.(size_memories) / log(2))
    lag_cat = conditional_to(observations, memories; past = past)

    n_in_t = [findall(n2t .== t) for t = 1:T]
    n_occurence_history = [findall(.&(observations[:, j] .== y, lag_cat[:, j] .== h)) for j = 1:D, h = 1:max_size_memory, y = 0:1]
    n_all = [n_per_category(tup..., n_in_t, n_occurence_history) for tup in Iterators.product(1:D, 1:max_size_memory, 1:T, 1:rain_cat)]

    model_A = model_for_A_d(s_ξ[:, 1, :], deg_Q, silence = silence) # JuMP Model for transition matrix
    model_B = model_for_B_d(γ_s[1, 1, 1, :, :], deg_Y, silence = silence) # JuMP Model for Emmission distribution

    loglikelihoods!(LL, hmm, observations, n2t, lag_cat)
    robust && replace!(LL, -Inf => nextfloat(-Inf), Inf => log(prevfloat(Inf)))

    forwardlog!(α, c, hmm.a, hmm.A, LL, n2t)
    backwardlog!(β, c, hmm.a, hmm.A, LL, n2t)
    posteriors!(γ, α, β)

    logtot = sum(c)
    (display == :iter) && println("Iteration 0: logtot = $logtot")

    for it = 1:maxiter
        update_a!(hmm.a, α, β)
        update_A_d!(hmm.A, θ_Q, ξ, s_ξ, α, β, LL, n2t, n_in_t, model_A; warm_start = warm_start)
        update_B_d!(hmm.B, θ_Y, γ, γ_s, observations, n_all, model_B; warm_start = warm_start)
        # Ensure the "connected-ness" of the states,
        # this prevents case where there is no transitions
        # between two extremely likely observations.
        robust && (hmm.A .+= eps())

        @check isprobvec(hmm.a)
        @check istransmats(hmm.A)

        push!(all_θ_Q, copy(θ_Q))
        push!(all_θ_Y, copy(θ_Y))

        # loglikelihoods!(LL, hmm, observations, n2t)
        loglikelihoods!(LL, hmm, observations, n2t, lag_cat)

        robust && replace!(LL, -Inf => nextfloat(-Inf), Inf => log(prevfloat(Inf)))

        forwardlog!(α, c, hmm.a, hmm.A, LL, n2t)
        backwardlog!(β, c, hmm.a, hmm.A, LL, n2t)
        posteriors!(γ, α, β)

        logtotp = sum(c)
        (display == :iter) && println(now(), " Iteration $it: logtot = $logtotp")
        flush(stdout)

        push!(history.logtots, logtotp)
        history.iterations += 1

        if abs(logtotp - logtot) < tol
            (display in [:iter, :final]) &&
                println("EM converged in $it iterations, logtot = $logtotp")
            history.converged = true
            break
        end

        logtot = logtotp
    end

    if !history.converged
        if display in [:iter, :final]
            println("EM has not converged after $(history.iterations) iterations, logtot = $logtot")
        end
    end

    history, all_θ_Q, all_θ_Y
end


text:
using LinearAlgebra

@testset "TwoFactorRepresentation - Basic Arithmetic" begin
    U1, Z1 = randn(100, 4), randn(1000, 4)

    U2, Z2 = randn(100, 2), randn(1000, 2)

    A = U1*Z1'
    lr_A = TwoFactorRepresentation(U1,Z1)

    B = U2*Z2'
    lr_B = TwoFactorRepresentation(U2,Z2)

    # testing addition
    @test norm(Matrix(lr_A + lr_B) - (A + B)) < 1e-10

    U3, Z3 = randn(1000, 7), randn(128, 7)
    C = U3*Z3'
    lr_C = TwoFactorRepresentation(U3, Z3)
    full_factor = randn(size(A))'

    # testing product
    @test norm(Matrix(lr_A*lr_C) - A*C) < 1e-10
    @test norm(Matrix(lr_A*full_factor) - A*full_factor) < 1e-10
    @test norm(Matrix(full_factor*lr_A) - full_factor*A) < 1e-10

    # testing elementwise product
    U4, Z4 = randn(100, 2), randn(1000,2)
    D = U4*Z4'
    lr_D = TwoFactorRepresentation(U4,Z4)
    @test norm(Matrix(lr_A .* lr_D) - A .* D) < 1e-10

    # testing elementwise power
    for d in 1:5
        @test norm(Matrix(lr_A .^ d) - A.^d) < 1e-8
    end

    # test sum
    lr_matrices = [TwoFactorRepresentation(randn(100, 5), randn(100, 5)) for i in 1:5]
    full_representations = [Matrix(lr_matrix) for lr_matrix in lr_matrices]
    @test norm(Matrix(sum(lr_matrices)) - sum(full_representations)) < 1e-8

    # test prod
    lr_matrices = [TwoFactorRepresentation(randn(100, 5), randn(100, 5)) for i in 1:5]
    full_representations = [Matrix(lr_matrix) for lr_matrix in lr_matrices]
    @test norm(Matrix(prod(lr_matrices)) -  prod(full_representations)) < 1e-5

    # indexing tests
    A = TwoFactorRepresentation([i + j for i in 1:10, j in 1:3], [i + j for i in 1:5, j in 1:3])
    # individual indexing
    for i in 1:10, j in 1:5
        @test A[i, j] == Matrix(A)[i,j]
    end
    # slicing
    for i in 1:10
        @test all(Matrix(A[i, :]) .== Matrix(A)[[i],:])
    end
    for j in 1:5
        @test all(Matrix(A[:, j]) .== Matrix(A)[:, [j]])
    end
    # more indexing
    @test all(Matrix(A[[1, 3, 7], [5,1]]) .== Matrix(A)[[1,3,7], [5,1]])
    @test all(Matrix(A[[1, 3, 7], 2]) .== Matrix(A)[[1,3,7], [2]])
    @test all(Matrix(A[2, [3,2]]) .== Matrix(A)[[2], [3,2]])

    # concatenation tests
    B = TwoFactorRepresentation([i + j for i in 1:10, j in 1:3], [i + j for i in 1:5, j in 1:3])
    @test all(Matrix(hcat(A,B)) .== hcat(Matrix(A), Matrix(B)))
    @test all(Matrix(vcat(A,B)) .== vcat(Matrix(A), Matrix(B)))
end

@testset "SVDLikeRepresentation - Basic Arithmetic" begin
    U1, S1, V1 = randn(100, 4), randn(4,4), randn(1000, 4)

    U2, S2, V2 = randn(100, 2), randn(2,2), randn(1000, 2)

    A = U1*S1*V1'
    lr_A = SVDLikeRepresentation(U1,S1,V1)

    B = U2*S2*V2'
    lr_B = SVDLikeRepresentation(U2,S2,V2)

    # testing addition
    @test norm(Matrix(lr_A + lr_B) - (A + B)) < 1e-10

    U3, S3, V3 = randn(1000, 7), randn(7,7), randn(128, 7)
    C = U3*S3*V3'
    lr_C = SVDLikeRepresentation(U3, S3, V3)
    full_factor = randn(size(A))'

    # testing product
    @test norm(Matrix(lr_A*lr_C) - A*C) < 1e-8
    @test norm(Matrix(lr_A*full_factor) - A*full_factor) < 1e-8
    @test norm(Matrix(full_factor*lr_A) - full_factor*A) < 1e-8

    # testing elementwise product
    U4, S4, V4 = randn(100, 2), randn(2,2), randn(1000,2)
    D = U4*S4*V4'
    lr_D = SVDLikeRepresentation(U4,S4,V4)
    @test norm(Matrix(lr_A .* lr_D) - A .* D) < 1e-8

    # testing elementwise power
    # very lenient test due to poor implementation
    for d in 1:5
        @test norm(Matrix(lr_A .^ d) - A.^d) < 1e-3
    end

    # test sum
    lr_matrices = [SVDLikeRepresentation(randn(100, 5), randn(5,5), randn(100, 5)) for i in 1:5]
    full_representations = [Matrix(lr_matrix) for lr_matrix in lr_matrices]
    @test norm(Matrix(sum(lr_matrices)) - sum(full_representations)) < 1e-8

    # test prod
    @test norm(Matrix(prod(lr_matrices)) -  prod(full_representations)) < 1e-5

    # indexing tests
    A = SVDLikeRepresentation([i + j for i in 1:10, j in 1:3], ones(3,3), [i + j for i in 1:5, j in 1:3])
    # individual indexing
    for i in 1:10, j in 1:5
        @test A[i, j] == Matrix(A)[i,j]
    end
    # slicing
    for i in 1:10
        @test all(Matrix(A[i, :]) .== Matrix(A)[[i],:])
    end
    for j in 1:5
        @test all(Matrix(A[:, j]) .== Matrix(A)[:, [j]])
    end
    # more indexing
    @test all(Matrix(A[[1, 3, 7], [5,1]]) .== Matrix(A)[[1,3,7], [5,1]])
    @test all(Matrix(A[[1, 3, 7], 2]) .== Matrix(A)[[1,3,7], [2]])
    @test all(Matrix(A[2, [3,2]]) .== Matrix(A)[[2], [3,2]])

    # concatenation tests
    B = SVDLikeRepresentation([i + j for i in 1:10, j in 1:3], ones(3,3), [i + j for i in 1:5, j in 1:3])
    @test all(Matrix(hcat(A,B)) .== hcat(Matrix(A), Matrix(B)))
    @test all(Matrix(vcat(A,B)) .== vcat(Matrix(A), Matrix(B)))
end


text:
"""
    trace_bath(rho, agg, FCProd, aggIndices, vibindices)

Trace out bath degrees of freedom from `rho`

``\\rho_\\text{tr} = \\operatorname{tr}_B \\{\\rho\\} =
\\sum_{k} \\langle k \\vert \\left( \\sum_{ab} \\rho_{am, bn} \\vert am \\rangle \\langle bn \\vert \\right)\\vert k \\rangle``

"""
function trace_bath_ground_excited(
    W::Array,
    aggCore::AggregateCore,
    aggTools::AggregateTools
)
    elLen = aggCore.molCount + 1
    rho = zeros(eltype(W), elLen, elLen)

    for I = 1:aggTools.bSize
        elind1, vibind1 = aggTools.indices[I]
        elOrder1 = OpenQuantumSystems.elIndOrder(elind1)

        for J = 1:aggTools.bSize
            elind2, vibind2 = aggTools.indices[J]
            elOrder2 = OpenQuantumSystems.elIndOrder(elind2)

            rho[elOrder1, elOrder2] += W[I, J] * aggTools.FCproduct[I, J]
        end
    end
    return rho
end

function trace_bath_ground_ground(
    W::Array,
    aggCore::AggregateCore,
    aggTools::AggregateTools
)
    elLen = aggCore.molCount + 1
    rho = zeros(eltype(W), elLen, elLen)
    indicesMap = aggTools.indicesMap

    for a=1:length(indicesMap)
        a1 = indicesMap[a][1]
        a2 = indicesMap[a][end]
        for b=1:length(indicesMap)
            b1 = indicesMap[b][1]
            b2 = indicesMap[b][end]
            rho[a, b] = tr(W[a1:a2, b1:b2])
        end
    end
    return rho
end

function trace_bath(
    W::Array,
    aggCore::AggregateCore,
    aggTools::AggregateTools;
    vib_basis::Symbol=:ground_excited
)
    if vib_basis ∉ (:ground_ground, :ground_excited)
        throw(ArgumentError("Optional argument vib_basis has to be selected from (:ground_ground, :ground_excited)"))
    end
    rho = nothing
    if vib_basis == :ground_excited
        rho = trace_bath_ground_excited(W, aggCore, aggTools)
    end
    if vib_basis == :ground_ground
        rho = trace_bath_ground_ground(W, aggCore, aggTools)
    end
    return rho
end

function trace_bath(
    W::Operator,
    aggCore::AggregateCore,
    aggTools::AggregateTools;
    vib_basis::Symbol=:ground_excited
)
    rho = trace_bath(W.data, aggCore, aggTools; vib_basis=vib_basis)
    return DenseOperator(aggTools.basisSystem, aggTools.basisSystem, rho)
end

"""
    trace_bath_slow(rho, agg, FCFact, aggIndices, vibindices)

Trace out bath degrees of freedom from `rho` without the product of Franck-Condon factors.

"""
function trace_bath_slow(
    W::Array,
    aggCore::AggregateCore,
    aggTools::AggregateTools
)
    elLen = aggCore.molCount + 1
    rho = zeros(eltype(W), elLen, elLen)

    for I = 1:aggTools.bSize
        elind1, vibind1 = aggTools.indices[I]
        elOrder1 = OpenQuantumSystems.elIndOrder(elind1)

        for J = 1:aggTools.bSize
            elind2, vibind2 = aggTools.indices[J]
            elOrder2 = OpenQuantumSystems.elIndOrder(elind2)

            for m = 1:aggTools.bBathSize
                # according to quantarhei, trace_over_vibrations()
                K = aggTools.indicesMap[1][m]
                L = aggTools.indicesMap[1][m]
                rho[elOrder1, elOrder2] += aggTools.FCfactors[K, I] * W[I, J] * aggTools.FCfactors[J, L]
            end
        end
    end
    return rho
end

function trace_bath_slow(
    W::Operator,
    aggCore::AggregateCore,
    aggTools::AggregateTools
)
    rho = trace_bath_slow(
        W.data,
        aggCore,
        aggTools
    )
    return DenseOperator(aggTools.basisSystem, aggTools.basisSystem, rho)
end

"""
    trace_bath(rho, a, b, agg, FCProd, aggIndices, vibindices)
Trace out bath degrees of freedom from `rho` without the product of Franck-Condon factors.
The trace will be done only on the Hilber space for electric bra part `a` and ket part `b`.
Input density matrix `rho` is for the whole Hilber space. This method returns number.
"""
function trace_bath(
    W::Array,
    a::N,
    b::N,
    aggTools::AggregateTools;
    vib_basis::Symbol=:ground_excited
) where {N <: Integer}
    if vib_basis ∉ (:ground_ground, :ground_excited)
        throw(ArgumentError("Optional argument vib_basis has to be selected from (:ground_ground, :ground_excited)"))
    end
    rho = eltype(W)(0)

    if vib_basis == :ground_excited
        for I in aggTools.indicesMap[a]
            for J in aggTools.indicesMap[b]
                rho += W[I, J] * aggTools.FCproduct[I, J]
            end
        end
    end
    if vib_basis == :ground_ground
        a1 = aggTools.indicesMap[a][1]
        a2 = aggTools.indicesMap[a][end]
        b1 = aggTools.indicesMap[b][1]
        b2 = aggTools.indicesMap[b][end]
        rho = tr(W[a1:a2, b1:b2])
    end

    rho
end

function trace_bath(
    W::Operator,
    a::N,
    b::N,
    aggTools::AggregateTools;
    vib_basis::Symbol=:ground_excited
) where {N <: Integer}
    rho = trace_bath(W.data, a, b, aggTools; vib_basis=vib_basis)
    return rho
end

"""
    trace_bath_part(rho, a, b, agg, FCProd, aggIndices, vibindices)
Trace out bath degrees of freedom from `rho` without the product of Franck-Condon factors.
The trace will be done only on the Hilber space for electric bra part `a` and ket part `b`.
Input density matrix `rho` is only for the subspace. This method returns number.
"""
function trace_bath_part(
    W::Array,
    a::N,
    b::N,
    aggTools::AggregateTools
) where {N <: Integer}
    rho_traced = eltype(W)(0)

    for a_vib = 1:aggTools.bBathSize
        I = aggTools.indicesMap[a][a_vib]
        for b_vib = 1:aggTools.bBathSize
            J = aggTools.indicesMap[b][b_vib]
            rho_traced += W[a_vib, b_vib] * aggTools.FCproduct[I, J]
        end
    end
    return rho_traced
end

function trace_bath_part(
    W::Operator,
    a::N,
    b::N,
    aggTools::AggregateTools
) where {N <: Integer}
    rho_traced = trace_bath_part(
        W.data,
        a,
        b,
        aggTools
    )
    return rho_traced
end

"""
    get_rho_bath(rho, agg, FCProd, aggIndices, vibindices; justCopy=false)

This method will return the bath part of `rho` knowing the result of [`trace_bath`](@ref) defined as follows

`` \\rho_\\text{bath} = \\operatorname{tr}_S \\{\\rho\\} ``

`` \\rho_{\\text{bath}, ab} = \\rho_{ab} / \\langle a \\vert \\operatorname{tr}_B \\{ \\rho \\}\\vert b \\rangle``

"""
function get_rho_bath(
    W::Array,
    aggCore::AggregateCore,
    aggTools::AggregateTools;
    justCopy::Bool = false,
    vib_basis::Symbol=:ground_excited
)
    rho = trace_bath(W, aggCore, aggTools; vib_basis=vib_basis)
    vibindices = aggTools.indicesMap
    vibLen = length(vibindices[end])
    elLen = aggCore.molCount + 1
    rho_bath = zeros(eltype(rho), aggTools.bSize, aggTools.bSize)
    rho_bath_ref = zeros(eltype(rho), vibLen, vibLen)

    el1_p = 0
    el2_p = 0
    for el1 = 1:elLen, el2 = 1:elLen
        if abs(rho[el1, el2]) != 0
            el1_p = el1
            el2_p = el2
            break
        end
    end
    vib11 = vibindices[el1_p][1]
    vib12 = vibindices[el1_p][end]
    vib21 = vibindices[el2_p][1]
    vib22 = vibindices[el2_p][end]
    rho_bath_ref[:, :] = W[vib11:vib12, vib21:vib22] / rho[el1_p, el2_p]
    for el1 = 1:elLen, el2 = 1:elLen
        vib11 = vibindices[el1][1]
        vib12 = vibindices[el1][end]
        vib21 = vibindices[el2][1]
        vib22 = vibindices[el2][end]
        if abs(rho[el1, el2]) != 0
            rho_bath[vib11:vib12, vib21:vib22] =
                W[vib11:vib12, vib21:vib22] / rho[el1, el2]
        else
            if justCopy
                rho_bath[vib11:vib12, vib21:vib22] = W[vib11:vib12, vib21:vib22]
            else
                rho_bath[vib11:vib12, vib21:vib22] = rho_bath_ref[:, :]
            end
        end
    end

    return rho_bath
end

function get_rho_bath(
    W::Operator,
    aggCore::AggregateCore,
    aggTools::AggregateTools;
    justCopy = false,
    vib_basis::Symbol=:ground_excited
)
    rho_data = get_rho_bath(
        W.data,
        aggCore,
        aggTools;
        justCopy = justCopy,
        vib_basis = vib_basis
    )
    return DenseOperator(W.basis_l, W.basis_r, rho_data)
end

"""
    ad(rho, rho_bath, agg, FCProd, aggIndices, vibindices)

This is the inverse operation to the trace over bath [`trace_bath`](@ref) and [`get_rho_bath`](@ref)
defined as follows

`` \\rho = \\operatorname{ad}\\{\\rho_\\text{tr}, \\rho_\\text{bath} \\} ``

"""
function ad(
    rho::Array,
    rho_bath::Array,
    aggCore::AggregateCore,
    aggTools::AggregateTools
)
    W = zero(rho_bath)

    for I = 1:aggTools.bSize
        elind1, vibind1 = aggTools.indices[I]
        elOrder1 = OpenQuantumSystems.elIndOrder(elind1)

        for J = 1:aggTools.bSize
            elind2, vibind2 = aggTools.indices[J]
            elOrder2 = OpenQuantumSystems.elIndOrder(elind2)

            W[I, J] = rho[elOrder1, elOrder2] * rho_bath[I, J]
        end
    end
    return W
end

function ad(
    rho::Operator,
    W_bath::Array,
    aggCore::AggregateCore,
    aggTools::AggregateTools;
)
    W = ad(
        rho.data,
        W_bath,
        aggCore,
        aggTools
    )
    return DenseOperator(aggTools.basis, aggTools.basis, W)
end

function ad(
    rho::Array,
    W_bath::Operator,
    aggCore::AggregateCore,
    aggTools::AggregateTools
)
    W = ad(
        rho,
        W_bath.data,
        aggCore,
        aggTools
    )
    return DenseOperator(aggTools.basis, aggTools.basis, W)
end

function ad(
    rho::Operator,
    W_bath::Operator,
    aggCore::AggregateCore,
    aggTools::AggregateTools;
)
    W = ad(
        rho.data,
        W_bath.data,
        aggCore,
        aggTools
    )
    return DenseOperator(aggTools.basis, aggTools.basis, W)
end

"""
    correlation_function(t, rho0_bath, Ham_0, Ham_I, agg, FCProd, aggInds, vibindices)

Get time dependent correlation function for a specified time `t` using following definition

`` C(t) = \\operatorname{tr}_B \\{ \\hat{H}_I^{(I)}(t) \\hat{H}_I \\rho_\\text{bath}(0) \\} ``

"""
function correlation_function(
    t,
    rho0_bath,
    Ham_0,
    Ham_I,
    aggCore::AggregateCore,
    aggTools::AggregateTools;
    vib_basis::Symbol=:ground_excited
)
    Ham_II_t = getInteractionHamIPicture(Ham_0, Ham_I, t)
    prod = Ham_II_t * Ham_I * rho0_bath
    return trace_bath(prod, aggCore, aggTools; vib_basis=vib_basis)
end

max_stars_repo_name: detrin/OpenQuantumSystems

max_issues_repo_name: detrin/OpenQuantumSystems

max_forks_repo_name: detrin/OpenQuantumSystems


text:
export BrickletHumidityHumidityCallbackThreshold
struct BrickletHumidityHumidityCallbackThreshold
    option::Char
    min::Integer
    max::Integer
end

export BrickletHumidityAnalogValueCallbackThreshold
struct BrickletHumidityAnalogValueCallbackThreshold
    option::Char
    min::Integer
    max::Integer
end

export BrickletHumidityIdentity
struct BrickletHumidityIdentity
    uid::String
    connected_uid::String
    position::Char
    hardware_version::Vector{Integer}
    firmware_version::Vector{Integer}
    device_identifier::Integer
end

export BrickletHumidity
"""
Measures relative humidity
"""
mutable struct BrickletHumidity <: TinkerforgeDevice
    ipcon::IPConnection
    deviceInternal::PyObject

    """
    Creates an object with the unique device ID *uid* and adds it to
    the IP Connection *ipcon*.
    """
    function BrickletHumidity(uid::String, ipcon::IPConnection)
        package = pyimport("tinkerforge.bricklet_humidity")
        deviceInternal = package.BrickletHumidity(uid, ipcon.ipconInternal)

        return new(ipcon, deviceInternal)
    end
end

export get_humidity
"""
    $(SIGNATURES)

Returns the humidity of the sensor.

If you want to get the humidity periodically, it is recommended to use the
:cb:`Humidity` callback and set the period with
:func:`Set Humidity Callback Period`.
"""
function get_humidity(device::BrickletHumidity)
    return device.deviceInternal.get_humidity()
end

export get_analog_value
"""
    $(SIGNATURES)

Returns the value as read by a 12-bit analog-to-digital converter.

.. note::
 The value returned by :func:`Get Humidity` is averaged over several samples
 to yield less noise, while :func:`Get Analog Value` gives back raw
 unfiltered analog values. The returned humidity value is calibrated for
 room temperatures, if you use the sensor in extreme cold or extreme
 warm environments, you might want to calculate the humidity from
 the analog value yourself. See the `HIH 5030 datasheet
 <https://github.com/Tinkerforge/humidity-bricklet/raw/master/datasheets/hih-5030.pdf>`__.

If you want the analog value periodically, it is recommended to use the
:cb:`Analog Value` callback and set the period with
:func:`Set Analog Value Callback Period`.
"""
function get_analog_value(device::BrickletHumidity)
    return device.deviceInternal.get_analog_value()
end

export set_humidity_callback_period
"""
    $(SIGNATURES)

Sets the period with which the :cb:`Humidity` callback is triggered
periodically. A value of 0 turns the callback off.

The :cb:`Humidity` callback is only triggered if the humidity has changed
since the last triggering.
"""
function set_humidity_callback_period(device::BrickletHumidity, period)
    device.deviceInternal.set_humidity_callback_period(period)
end

export get_humidity_callback_period
"""
    $(SIGNATURES)

Returns the period as set by :func:`Set Humidity Callback Period`.
"""
function get_humidity_callback_period(device::BrickletHumidity)
    return device.deviceInternal.get_humidity_callback_period()
end

export set_analog_value_callback_period
"""
    $(SIGNATURES)

Sets the period with which the :cb:`Analog Value` callback is triggered
periodically. A value of 0 turns the callback off.

The :cb:`Analog Value` callback is only triggered if the analog value has
changed since the last triggering.
"""
function set_analog_value_callback_period(device::BrickletHumidity, period)
    device.deviceInternal.set_analog_value_callback_period(period)
end

export get_analog_value_callback_period
"""
    $(SIGNATURES)

Returns the period as set by :func:`Set Analog Value Callback Period`.
"""
function get_analog_value_callback_period(device::BrickletHumidity)
    return device.deviceInternal.get_analog_value_callback_period()
end

export set_humidity_callback_threshold
"""
    $(SIGNATURES)

Sets the thresholds for the :cb:`Humidity Reached` callback.

The following options are possible:

.. csv-table::
 :header: "Option", "Description"
 :widths: 10, 100

 "'x'",    "Callback is turned off"
 "'o'",    "Callback is triggered when the humidity is *outside* the min and max values"
 "'i'",    "Callback is triggered when the humidity is *inside* the min and max values"
 "'<'",    "Callback is triggered when the humidity is smaller than the min value (max is ignored)"
 "'>'",    "Callback is triggered when the humidity is greater than the min value (max is ignored)"
"""
function set_humidity_callback_threshold(device::BrickletHumidity, option, min, max)
    device.deviceInternal.set_humidity_callback_threshold(option, min, max)
end

export get_humidity_callback_threshold
"""
    $(SIGNATURES)

Returns the threshold as set by :func:`Set Humidity Callback Threshold`.
"""
function get_humidity_callback_threshold(device::BrickletHumidity)
    return device.deviceInternal.get_humidity_callback_threshold()
end

export set_analog_value_callback_threshold
"""
    $(SIGNATURES)

Sets the thresholds for the :cb:`Analog Value Reached` callback.

The following options are possible:

.. csv-table::
 :header: "Option", "Description"
 :widths: 10, 100

 "'x'",    "Callback is turned off"
 "'o'",    "Callback is triggered when the analog value is *outside* the min and max values"
 "'i'",    "Callback is triggered when the analog value is *inside* the min and max values"
 "'<'",    "Callback is triggered when the analog value is smaller than the min value (max is ignored)"
 "'>'",    "Callback is triggered when the analog value is greater than the min value (max is ignored)"
"""
function set_analog_value_callback_threshold(device::BrickletHumidity, option, min, max)
    device.deviceInternal.set_analog_value_callback_threshold(option, min, max)
end

export get_analog_value_callback_threshold
"""
    $(SIGNATURES)

Returns the threshold as set by :func:`Set Analog Value Callback Threshold`.
"""
function get_analog_value_callback_threshold(device::BrickletHumidity)
    return device.deviceInternal.get_analog_value_callback_threshold()
end

export set_debounce_period
"""
    $(SIGNATURES)

Sets the period with which the threshold callbacks

* :cb:`Humidity Reached`,
* :cb:`Analog Value Reached`

are triggered, if the thresholds

* :func:`Set Humidity Callback Threshold`,
* :func:`Set Analog Value Callback Threshold`

keep being reached.
"""
function set_debounce_period(device::BrickletHumidity, debounce)
    device.deviceInternal.set_debounce_period(debounce)
end

export get_debounce_period
"""
    $(SIGNATURES)

Returns the debounce period as set by :func:`Set Debounce Period`.
"""
function get_debounce_period(device::BrickletHumidity)
    return device.deviceInternal.get_debounce_period()
end

export get_identity
"""
    $(SIGNATURES)

Returns the UID, the UID where the Bricklet is connected to,
the position, the hardware and firmware version as well as the
device identifier.

The position can be 'a', 'b', 'c', 'd', 'e', 'f', 'g' or 'h' (Bricklet Port).
A Bricklet connected to an :ref:`Isolator Bricklet <isolator_bricklet>` is always at
position 'z'.

The device identifier numbers can be found :ref:`here <device_identifier>`.
|device_identifier_constant|
"""
function get_identity(device::BrickletHumidity)
    return device.deviceInternal.get_identity()
end

export register_callback
"""
Registers the given *function* with the given *callback_id*.
"""
function register_callback(device::BrickletHumidity, callback_id, function_)
    if isnothing(function_)
        device.registered_callbacks.pop(callback_id, None)
    else
        device.registered_callbacks[callback_id] = function_
    end
end