using OrdinaryDiffEq, SparseArrays, DelimitedFiles, FastLapackInterface
using SpecialFunctions: besselj0, besselj
using LinearAlgebra: LAPACK, BLAS, Eigen, Symmetric, I, diagind, diag, eigen, eigvals, mul!
using Distributed, SharedArrays
using ProgressMeter: @showprogress
import ProgressMeter
using FLoops: @floop, @init
import FLoops
import Base.show, Base.copy

"""
A type representing a Bose-Hubbard Hamiltonian,
    H = - âˆ‘ ğ��½áµ¢â±¼ ğ�‘�â€ áµ¢ ğ�‘�â±¼, + ğ�‘ˆ/2 âˆ‘ ğ�‘›áµ¢(ğ�‘›áµ¢ - 1).
"""
mutable struct BoseHamiltonian{Float <: AbstractFloat}
    lattice::Lattice
    J::Float
    U::Float
    f::Float # F / Ï‰
    Ï‰::Float
    Ï‰â‚—::Float # lower bound of the Floquet zone; needed for (E)DPT calculations
    r::Rational{Int} # resonance number; needed for DPT calculations
    type::Symbol     # `:edpt`, `:dpt`, `:diverging` or anything else for ordinary high-frequency expansion
    order::Int  # oder of (E)DPT
    space_of_state::Vector{Tuple{Int,Int}}    # space_of_state[i] stores the subspace number (ğ��´, ğ�‘�) of i'th state, with ğ��´ = 0 assigned to all nondegenerate space
    H::Matrix{Float}   # the Hamiltonian matrix
    Eâ‚€::Vector{Int}    # zeroth-order spectrum, in units of ğ�‘ˆ
    Îµâ‚€::Vector{Float}  # zeroth-order quasienergy spectrum
    R1::Dict{Tuple{Int,Int,Int,Int,Int,Int,Int,Bool}, Float} # required for (E)DPT calculation
    R2::Dict{Tuple{Int,Int,Int,Int,Int,Int,Int,Int,Int,Int,Bool}, Float} # required for (E)DPT-3 calculation
    R3::Dict{NTuple{7, Int64}, Float} # required for (E)DPT-3 calculation
end

"""
Construct a `BoseHamiltonian` object defined on `lattice`.
Type of `J` determines the type of Float used for all fields of the resulting object.
`Ï‰â‚—` is the lower bound of the first Floquet zone.
"""
function BoseHamiltonian(lattice::Lattice, J::Float, U::Real, f::Real, Ï‰::Real; Ï‰â‚—::Real=-Ï‰/2, r::Rational=0//1, order::Integer=1, type::Symbol=:basic) where {Float <: AbstractFloat}
    nstates = length(lattice.basis_states)
    # Calculate zeroth-order spectrum, in units of ğ�‘ˆ. It only depends on the lattice, so will not change
    Eâ‚€ = zeros(Int, nstates)
    for (index, state) in enumerate(lattice.basis_states)
        for n_i in state
            if (n_i > 1)
                Eâ‚€[index] += n_i * (n_i - 1) Ã· 2 # will always be divisible by 2
            end
        end
    end
    H = Matrix{Float}(undef, nstates, nstates)
    Îµâ‚€ = Vector{Float}(undef, nstates)
    R1 = Dict{Tuple{Int,Int,Int,Int,Int,Int,Int,Bool}, Float}()
    R2 = Dict{Tuple{Int,Int,Int,Int,Int,Int,Int,Int,Int,Int,Bool}, Float}()
    R3 = Dict{NTuple{7, Int64}, Float}()
    space_of_state = (type in (:edpt, :dpt) ? Vector{Tuple{Int,Int}}(undef, nstates) : Vector{Tuple{Int,Int}}())

    bh = BoseHamiltonian(lattice, Float(J), Float(U), Float(f), Float(Ï‰), Float(Ï‰â‚—), r, type, order, space_of_state, H, Eâ‚€, Îµâ‚€, R1, R2, R3)
    update_params!(bh)

    return bh
end

"Print non-zero elements of the Hamiltonian `bh` in the format âŸ¨bra| HÌ‚ |ketâŸ©."
function Base.show(io::IO, bh::BoseHamiltonian{<:AbstractFloat})
    for C in CartesianIndices(bh.H)
        bh.H[C[1], C[2]] != 0 && println(io, bh.lattice.basis_states[C[1]], " HÌ‚ ", bh.lattice.basis_states[C[2]], " = ", round(bh.H[C[1], C[2]], sigdigits=3))
    end
end

"""
Return a shallow copy of a given `BoseHamiltonian`: copy all scalar fields; reference the same `bh.lattice` and hence copy `bh.Eâ‚€`;
make similar `bh.Îµâ‚€`; make empty copies of `bh.R*` so that they are not populated in the resulting object but are of the same allocated size (think `sizehint`).
"""
function Base.copy(bh::BoseHamiltonian{Float}) where {Float<:AbstractFloat}
    # since `Eâ‚€` depends on the lattice, we are copying `bh.Eâ‚€`
    b = BoseHamiltonian(bh.lattice, bh.J, bh.U, bh.f, bh.Ï‰, bh.Ï‰â‚—, bh.r, bh.type, bh.order, similar(bh.space_of_state), similar(bh.H),
        copy(bh.Eâ‚€), similar(bh.Îµâ‚€), empty(bh.R1), empty(bh.R2), empty(bh.R3))
    sizehint!(b.R1, length(bh.R1)) 
    sizehint!(b.R2, length(bh.R2)) 
    sizehint!(b.R3, length(bh.R3))
    return b 
end

"Update parameters of `bh` and reconstruct `bh.H`."
function update_params!(bh::BoseHamiltonian{<:AbstractFloat}; J::Real=bh.J, U::Real=bh.U, f::Real=bh.f, Ï‰::Real=bh.Ï‰, Ï‰â‚—::Real=bh.Ï‰â‚—, r::Rational=bh.r, order::Integer=bh.order, type::Symbol=bh.type)
    bh.J = J; bh.U = U; bh.f = f; bh.Ï‰ = Ï‰; bh.Ï‰â‚— = Ï‰â‚—; bh.r = r; bh.order = order; bh.type = type
    if type == :edpt
        map!(bh.space_of_state, bh.Eâ‚€) do E
            # rounding helps in cases such as when E*U - Ï‰â‚— = 29.9...96 and Ã·10 gives 2 instead of 3
            a = round(E*U - Ï‰â‚—, sigdigits=6) Ã· Ï‰ |> Int
            A = E % denominator(r) # not used in EDPT!
            return (A, a)
        end
        for i in eachindex(bh.Eâ‚€)
            bh.Îµâ‚€[i] = bh.Eâ‚€[i]*U - bh.space_of_state[i][2]*Ï‰
        end
        constructH_edpt!(bh, order)
    elseif type == :dpt
        map!(bh.space_of_state, bh.Eâ‚€) do E
            a = round(E*r*Ï‰ - Ï‰â‚—, sigdigits=6) Ã· Ï‰ |> Int
            A = clamp(E % denominator(r), 0, 1) # if division result is 0, then A = 0, otherwise A = 1
            return (A, a)
        end
        for i in eachindex(bh.Eâ‚€)
            bh.Îµâ‚€[i] = bh.Eâ‚€[i]*U - bh.space_of_state[i][2]*Ï‰
        end
        constructH_dpt!(bh, order)
    elseif type == :diverging
        constructH_diverging!(bh, order)
    else
        constructH!(bh, order)
    end
end

"Construct the Hamiltonian matrix."
function constructH!(bh::BoseHamiltonian{Float}, order::Integer) where {Float<:AbstractFloat}
    (;J, U, f, Ï‰, H) = bh
    (;index_of_state, ncells, neis_of_cell) = bh.lattice
    H .= 0
    H[diagind(H)] .= bh.Eâ‚€ .* U

    Jeff = J * besselj0(f)

    J_sum = zeros(Float, 2)
    if order == 2
        a_max = 20
        J_sum[1] = (J/Ï‰)^2 * U * sum(         besselj(a, f)^2 / a^2 for a in 1:a_max) # for ğ�‘— = k
        J_sum[2] = (J/Ï‰)^2 * U * sum((-1)^a * besselj(a, f)^2 / a^2 for a in 1:a_max) # for ğ�‘— â‰  k
    end

    # take each basis state and find which transitions are possible
    for (ket, Î±) in index_of_state
        for i = 1:ncells # iterate over the terms of the Hamiltonian
            # ğ�‘�â€ áµ¢ ğ�‘�â±¼
            for (j, _) in neis_of_cell[i]
                if (ket[j] > 0) # check that a particle is present at site `j` so that destruction ğ�‘�â±¼ is possible
                    val = -Jeff * sqrt( (ket[i]+1) * ket[j] )
                    bra = copy(ket)
                    bra[j] -= 1
                    bra[i] += 1
                    Î±â€² = index_of_state[bra]
                    H[Î±â€², Î±] += val
                end
            end

            if order == 2
                for (j, _) in neis_of_cell[i], (k, _) in neis_of_cell[i]
                    Câ‚�, Câ‚‚ = j == k ? J_sum : (J_sum[2], J_sum[1])
                    # ğ�‘�â€ â‚– (2ğ�‘›áµ¢ - ğ�‘›â±¼) ğ�‘�â±¼
                    if (ket[j] > 0 && ket[i] != ket[j]-1)
                        val = Câ‚� * âˆš( (k == j ? ket[k] : ket[k]+1) * ket[j] ) * (2ket[i] - (ket[j]-1))
                        bra = copy(ket)
                        bra[j] -= 1
                        bra[k] += 1
                        Î±â€² = index_of_state[bra]
                        H[Î±â€², Î±] += val
                    end
                    # ğ�‘�â€ â±¼ (2ğ�‘›áµ¢ - ğ�‘›â±¼) ğ�‘�â‚–
                    if (ket[k] > 0 && ket[i] != (j == k ? ket[j]-1 : ket[j]))
                        val = Câ‚� * âˆš( (j == k ? ket[j] : ket[j]+1) * ket[k] ) * (2ket[i] - (j == k ? ket[j]-1 : ket[j]))
                        bra = copy(ket)
                        bra[k] -= 1
                        bra[j] += 1
                        Î±â€² = index_of_state[bra]
                        H[Î±â€², Î±] += val
                    end
                    # ğ�‘�â€ áµ¢ ğ�‘�â€ áµ¢ ğ�‘�â‚– ğ�‘�â±¼
                    if ( (k == j && ket[j] > 1) || (k != j && ket[k] > 0 && ket[j] > 0))
                        val = -Câ‚‚ * âˆš( (ket[i]+2) * (ket[i]+1) * (k == j ? (ket[j]-1)ket[j] : ket[j]ket[k]))
                        bra = copy(ket)
                        bra[j] -= 1
                        bra[k] -= 1
                        bra[i] += 2
                        Î±â€² = index_of_state[bra]
                        H[Î±â€², Î±] += val
                    end
                    # ğ�‘�â€ â‚– ğ�‘�â€ â±¼ ğ�‘�áµ¢ ğ�‘�áµ¢
                    if (ket[i] > 1)
                        val = -Câ‚‚ * âˆš( (k == j ? (ket[j]+2) * (ket[j]+1) : (ket[k]+1) * (ket[j]+1)) * (ket[i]-1)ket[i])
                        bra = copy(ket)
                        bra[i] -= 2
                        bra[j] += 1
                        bra[k] += 1
                        Î±â€² = index_of_state[bra]
                        H[Î±â€², Î±] += val
                    end
                end
            end
        end
    end
end

"""
Construct the Hamiltonian matrix for the degenerate case but without (E)DPT.
We do not assume that ğ�‘ˆ â‰ª ğ�œ”, but we do not use (E)DPT either, leading to diverging results.
"""
function constructH_diverging!(bh::BoseHamiltonian{Float}, order::Integer) where {Float<:AbstractFloat}
    (;J, U, f, Ï‰, H) = bh
    (;index_of_state, ncells, nbozons, neis_of_cell) = bh.lattice
    H .= 0
    H[diagind(H)] .= bh.Eâ‚€ .* U

    Jeff = J * besselj0(f)

    n_max = nbozons - 1
    n_min = -nbozons - 1
    R1 = Dict{Int64, Real}()
    R2 = Dict{Int64, Real}()
    for n in n_min:n_max
        R1[n] = ğ�‘…(Ï‰, U*n, f, type=1)
        R2[n] = ğ�‘…(Ï‰, U*n, f, type=2)
    end

    # take each basis state and find which transitions are possible
    for (ket, Î±) in index_of_state
        for i = 1:ncells # iterate over the terms of the Hamiltonian
            # ğ�‘�â€ áµ¢ ğ�‘�â±¼
            for (j, _) in neis_of_cell[i]
                if (ket[j] > 0) # check that a particle is present at site `j` so that destruction ğ�‘�â±¼ is possible
                    val = -Jeff * sqrt( (ket[i]+1) * ket[j] )
                    bra = copy(ket)
                    bra[j] -= 1
                    bra[i] += 1
                    Î±â€² = index_of_state[bra]
                    H[Î±â€², Î±] += val
                end
            end

            if order == 2
                for (j, i_j) in neis_of_cell[i], k in 1:ncells, (l, k_l) in neis_of_cell[k]
                    # ğ�‘�â€ áµ¢ ğ�‘�â±¼ [ğ�‘�ğ�œ”+ğ�‘ˆ(ğ�‘›â‚–-ğ�‘›â‚—-1)]â�»Â¹ ğ�‘�â€ â‚– ğ�‘�â‚—
                    if ( ket[l] > 0 && (j == k || (j == l && ket[j] > 1) || (j != l && ket[j] > 0)) )
                        R = i_j == k_l ? R1 : R2
                        val = -J^2/2
                        bra = copy(ket)
                        val *= âˆšbra[l]
                        bra[l] -= 1
                        bra[k] += 1
                        val *= R[bra[k] - bra[l] - 1] * âˆšbra[k]
                        val *= âˆšbra[j]
                        bra[j] -= 1
                        bra[i] += 1
                        val *= âˆšbra[i]
                        Î±â€² = index_of_state[bra]
                        H[Î±â€², Î±] += val
                    end

                    # [ğ�‘�ğ�œ”+ğ�‘ˆ(ğ�‘›â‚–-ğ�‘›â‚—-1)]â�»Â¹ ğ�‘�â€ â‚– ğ�‘�â‚— ğ�‘�â€ áµ¢ ğ�‘�â±¼ 
                    if ( ket[j] > 0 && (l == i || (l == j && ket[l] > 1) || (l != j && ket[l] > 0)) )
                        R = i_j == k_l ? R1 : R2
                        val = +J^2/2
                        bra = copy(ket)
                        val *= âˆšbra[j]
                        bra[j] -= 1
                        bra[i] += 1
                        val *= âˆšbra[i]
                        val *= âˆšbra[l]
                        bra[l] -= 1
                        bra[k] += 1
                        val *= R[bra[k] - bra[l] - 1] * âˆšbra[k]
                        Î±â€² = index_of_state[bra]
                        H[Î±â€², Î±] += val
                    end
                end
            end
        end
    end
end

function ğ�‘…(Ï‰::Real, Un::Real, f::Real; type::Integer)
    N = 20
    aâ‚€ = round(Int, -Un / Ï‰)
    # if `Un / Ï‰` is integer, aâ‚€ should be skipped in the sum
    a_range = isinteger(Un / Ï‰) ? [aâ‚€-N:aâ‚€-1; aâ‚€+1:aâ‚€+N] : collect(aâ‚€-N:aâ‚€+N) # collect for type stability
    r = zero(Ï‰)
    if type == 1
        for a in a_range
            a == 0 && continue
            r += 1/(a*Ï‰ + Un) * besselj(a, f)^2 * (-1)^a
        end
    else
        for a in a_range
            a == 0 && continue
            r += 1/(a*Ï‰ + Un) * besselj(a, f)^2
        end
    end
    return r
end

"Construct the Hamiltonian matrix using EDPT."
function constructH_edpt!(bh::BoseHamiltonian{Float}, order::Integer) where {Float<:AbstractFloat}
    (;index_of_state, ncells, neis_of_cell) = bh.lattice
    (;J, U, f, Ï‰, Eâ‚€, Îµâ‚€, space_of_state, H, R1, R2, R3) = bh

    H .= 0
    H[diagind(H)] .= Îµâ‚€

    empty!(R1); empty!(R2); empty!(R3);

    bra = similar(bh.lattice.basis_states[1])
    # take each basis state and find which transitions are possible
    for (ket, Î±) in index_of_state
        A, a = space_of_state[Î±]
        for i = 1:ncells # iterate over the terms of the Hamiltonian
            # 1st order
            for (j, i_j) in neis_of_cell[i]
                if (ket[j] > 0) # check that a particle is present at site `j` so that destruction ğ�‘�â±¼ is possible
                    copy!(bra, ket)
                    bra[j] -= 1
                    bra[i] += 1
                    Î±â€² = index_of_state[bra]
                    _, aâ€² = space_of_state[Î±â€²]
                    val = -J * besselj(a - aâ€², f*i_j) * sqrt( (ket[i]+1) * ket[j] )
                    H[Î±â€², Î±] += val
                end
            end

            if order >= 2
                for (j, i_j) in neis_of_cell[i]
                    for k = 1:ncells, (l, k_l) in neis_of_cell[k]
                        # ğ�‘�â€ áµ¢ ğ�‘�â±¼ ğ�‘�â€ â‚– ğ�‘�â‚—
                        if ( ket[l] > 0 && (j == k || (j == l && ket[j] > 1) || (j != l && ket[j] > 0)) )
                            val = +J^2/2
                            copy!(bra, ket)
                            val *= âˆšbra[l]
                            bra[l] -= 1
                            bra[k] += 1
                            val *= âˆšbra[k]
                            _, b = space_of_state[index_of_state[bra]]
                            val *= âˆšbra[j]
                            bra[j] -= 1
                            bra[i] += 1
                            Î±â€² = index_of_state[bra]
                            _, aâ€² = space_of_state[Î±â€²]
                            val *= âˆšbra[i]
                            val *= (get_R!(R1, U, Ï‰, f, bra[i]-bra[j]-1, aâ€²-b, i_j, k_l, aâ€², a, b, true) +
                                    get_R!(R1, U, Ï‰, f, ket[l]-ket[k]-1, a-b, i_j, k_l, aâ€², a, b, true))
                            H[Î±â€², Î±] += val
                        end
                    end
                end
            end

            if order >= 3
                for (j, i_j) in neis_of_cell[i]
                    for k = 1:ncells, (l, k_l) in neis_of_cell[k], m = 1:ncells, (n, m_n) in neis_of_cell[m]
                        # ğ�‘�â€ áµ¢ ğ�‘�â±¼ ğ�‘�â€ â‚– ğ�‘�â‚— ğ�‘�â€ â‚˜ ğ�‘�â‚™
                        ket[n] == 0 && continue
                        copy!(bra, ket)
                        val = -J^3/2
                        val *= âˆšbra[n]; bra[n] -= 1; bra[m] += 1; val *= âˆšbra[m]
                        bra[l] == 0 && continue
                        Î³ = index_of_state[bra]
                        _, c = space_of_state[Î³]
                        val *= âˆšbra[l]; bra[l] -= 1; bra[k] += 1; val *= âˆšbra[k]
                        bra[j] == 0 && continue
                        Î² = index_of_state[bra]
                        _, b = space_of_state[Î²]
                        val *= âˆšbra[j]; bra[j] -= 1; bra[i] += 1; val *= âˆšbra[i]
                        Î±â€² = index_of_state[bra]
                        _, aâ€² = space_of_state[Î±â€²]

                        s = zero(Float) # terms of the sum
                        J_indices = (-aâ€²+b, -b+c, -c+a)
                        J_args = (i_j, k_l, m_n)
                        Î”E1 = Eâ‚€[Î³] - Eâ‚€[Î±â€²]
                        Î”E2 = Eâ‚€[Î²] - Eâ‚€[Î³]
                        s += get_R2!(R2, U, Ï‰, f, Î”E1, Î”E2, c-aâ€², b-c, J_indices, J_args, true)
                    
                        J_indices = (a-c, b-aâ€², c-b)
                        J_args = (m_n, i_j, k_l)
                        Î”E1 = Eâ‚€[Î²] - Eâ‚€[Î±]
                        Î”E2 = Eâ‚€[Î³] - Eâ‚€[Î²]
                        s += get_R2!(R2, U, Ï‰, f, Î”E1, Î”E2, b-a, c-b, J_indices, J_args, true)
                    
                        J_indices = (c-b, b-aâ€², a-c)
                        J_args = (k_l, i_j, m_n)
                        Î”E1 = Eâ‚€[Î²] - Eâ‚€[Î±â€²]
                        Î”E2 = Eâ‚€[Î±â€²] - Eâ‚€[Î³]
                        s -= get_R2!(R2, U, Ï‰, f, Î”E1, Î”E2, b-aâ€², aâ€²-c, J_indices, J_args, true)

                        Î”E1 = Eâ‚€[Î²] - Eâ‚€[Î±]
                        Î”E2 = Eâ‚€[Î±] - Eâ‚€[Î³]
                        s -= get_R2!(R2, U, Ï‰, f, Î”E1, Î”E2, b-a, a-c, J_indices, J_args, true)

                        key = (Eâ‚€[Î±], Eâ‚€[Î²], Eâ‚€[Î³], Eâ‚€[Î±â€²], i_j, k_l, m_n)
                        N = 20
                        s += get!(R3, key) do
                            t = zero(Float)
                            for p in -N:N
                                p == 0 && continue
                                for q in -N:N
                                    (q == 0 || q == p) && continue
                                    # cast to Int32 as otherwise multiplication of three `besselj`s leads to allocations if `typeof(f) == Float32`
                                    t += besselj(Int32(b-aâ€²-p), f*i_j) * besselj(Int32(c-b+p-q), f*k_l) * besselj(Int32(a-c+q), f*m_n) * (
                                            1 / 2(Îµâ‚€[Î±â€²] - Îµâ‚€[Î³] - q*Ï‰)     * (1/(Îµâ‚€[Î³] - Îµâ‚€[Î²]  - (p-q)*Ï‰) - 1/(Îµâ‚€[Î²] - Îµâ‚€[Î±â€²] + p*Ï‰)) +
                                            1 / 2(Îµâ‚€[Î±]  - Îµâ‚€[Î²] - p*Ï‰)     * (1/(Îµâ‚€[Î±] - Îµâ‚€[Î³]  - q*Ï‰)     - 1/(Îµâ‚€[Î³] - Îµâ‚€[Î²]  - (p-q)*Ï‰)) +
                                            1 / 6(Îµâ‚€[Î³]  - Îµâ‚€[Î²] - (p-q)*Ï‰) * (1/(Îµâ‚€[Î²] - Îµâ‚€[Î±â€²] + p*Ï‰)     + 1/(Îµâ‚€[Î±] - Îµâ‚€[Î³]  - q*Ï‰)) -
                                            1 / 3(Îµâ‚€[Î±]  - Îµâ‚€[Î³] - q*Ï‰) / (Îµâ‚€[Î²] - Îµâ‚€[Î±â€²] + p*Ï‰) )
                                end
                            end
                            t
                        end
                        val *= s
                        H[Î±â€², Î±] += val
                    end
                end
            end
        end
    end
end

"Construct the Hamiltonian matrix using DPT."
function constructH_dpt!(bh::BoseHamiltonian{Float}, order::Integer) where {Float<:AbstractFloat}
    (;index_of_state, ncells, neis_of_cell) = bh.lattice
    (;J, U, f, Ï‰, Eâ‚€, Îµâ‚€, space_of_state, H, R1, R2, R3) = bh

    H .= 0
    H[diagind(H)] .= Îµâ‚€

    empty!(R1); empty!(R2); empty!(R3);
    
    bra = similar(bh.lattice.basis_states[1])
    # take each basis state and find which transitions are possible
    for (ket, Î±) in index_of_state
        A, a = space_of_state[Î±]
        for i = 1:ncells # iterate over the terms of the Hamiltonian
            # 1st order
            for (j, i_j) in neis_of_cell[i]
                if (ket[j] > 0) # check that a particle is present at site `j` so that destruction ğ�‘�â±¼ is possible
                    copy!(bra, ket)
                    bra[j] -= 1
                    bra[i] += 1
                    Î±â€² = index_of_state[bra]
                    Aâ€², aâ€² = space_of_state[Î±â€²]
                    if Aâ€² == A # proceed only if bra is in the same degenerate space
                        val = -J * besselj(a - aâ€², f*i_j) * sqrt( (ket[i]+1) * ket[j] )
                        H[Î±â€², Î±] += val
                    end
                end
            end

            if order >= 2
                for (j, i_j) in neis_of_cell[i]
                    for k = 1:ncells, (l, k_l) in neis_of_cell[k]
                        # ğ�‘�â€ áµ¢ ğ�‘�â±¼ ğ�‘�â€ â‚– ğ�‘�â‚—
                        if ( ket[l] > 0 && (j == k || (j == l && ket[j] > 1) || (j != l && ket[j] > 0)) )
                            val = +J^2/2
                            copy!(bra, ket)
                            val *= âˆšbra[l]
                            bra[l] -= 1
                            bra[k] += 1
                            val *= âˆšbra[k]
                            B, b = space_of_state[index_of_state[bra]]
                            val *= âˆšbra[j]
                            bra[j] -= 1
                            bra[i] += 1
                            Î±â€² = index_of_state[bra]
                            Aâ€², aâ€² = space_of_state[Î±â€²]
                            if Aâ€² == A # proceed only if bra is in the same degenerate space
                                val *= âˆšbra[i]
                                skipzero = (B == A)
                                val *= (get_R!(R1, U, Ï‰, f, bra[i]-bra[j]-1, aâ€²-b, i_j, k_l, aâ€², a, b, skipzero) +
                                        get_R!(R1, U, Ï‰, f, ket[l]-ket[k]-1, a-b, i_j, k_l, aâ€², a, b, skipzero))
                                H[Î±â€², Î±] += val
                            end
                        end
                    end
                end
            end

            if order >= 3
                for (j, i_j) in neis_of_cell[i]
                    for k = 1:ncells, (l, k_l) in neis_of_cell[k], m = 1:ncells, (n, m_n) in neis_of_cell[m]
                        # ğ�‘�â€ áµ¢ ğ�‘�â±¼ ğ�‘�â€ â‚– ğ�‘�â‚— ğ�‘�â€ â‚˜ ğ�‘�â‚™
                        ket[n] == 0 && continue
                        copy!(bra, ket)
                        val = -J^3/2
                        val *= âˆšbra[n]; bra[n] -= 1; bra[m] += 1; val *= âˆšbra[m]
                        bra[l] == 0 && continue
                        Î³ = index_of_state[bra]
                        C, c = space_of_state[Î³]
                        val *= âˆšbra[l]; bra[l] -= 1; bra[k] += 1; val *= âˆšbra[k]
                        bra[j] == 0 && continue
                        Î² = index_of_state[bra]
                        B, b = space_of_state[Î²]
                        val *= âˆšbra[j]; bra[j] -= 1; bra[i] += 1; val *= âˆšbra[i]
                        Î±â€² = index_of_state[bra]
                        Aâ€², aâ€² = space_of_state[Î±â€²]
                        if Aâ€² == A
                            s = zero(Float) # terms of the sum
                            if A == B
                                J_indices = (-aâ€²+b, -b+c, -c+a)
                                J_args = (i_j, k_l, m_n)
                                Î”E1 = Eâ‚€[Î³] - Eâ‚€[Î±â€²]
                                Î”E2 = Eâ‚€[Î²] - Eâ‚€[Î³]
                                s += get_R2!(R2, U, Ï‰, f, Î”E1, Î”E2, c-aâ€², b-c, J_indices, J_args, B == C)
                            end
                            if A == C
                                J_indices = (a-c, b-aâ€², c-b)
                                J_args = (m_n, i_j, k_l)
                                Î”E1 = Eâ‚€[Î²] - Eâ‚€[Î±]
                                Î”E2 = Eâ‚€[Î³] - Eâ‚€[Î²]
                                s += get_R2!(R2, U, Ï‰, f, Î”E1, Î”E2, b-a, c-b, J_indices, J_args, B == C)
                            end
                            if B == C
                                J_indices = (c-b, b-aâ€², a-c)
                                J_args = (k_l, i_j, m_n)
                                Î”E1 = Eâ‚€[Î²] - Eâ‚€[Î±â€²]
                                Î”E2 = Eâ‚€[Î±â€²] - Eâ‚€[Î³]
                                s -= get_R2!(R2, U, Ï‰, f, Î”E1, Î”E2, b-aâ€², aâ€²-c, J_indices, J_args, B == C)

                                Î”E1 = Eâ‚€[Î²] - Eâ‚€[Î±]
                                Î”E2 = Eâ‚€[Î±] - Eâ‚€[Î³]
                                s -= get_R2!(R2, U, Ï‰, f, Î”E1, Î”E2, b-a, a-c, J_indices, J_args, B == C)
                            end

                            key = (Eâ‚€[Î±], Eâ‚€[Î²], Eâ‚€[Î³], Eâ‚€[Î±â€²], i_j, k_l, m_n)
                            N = 20
                            s += get!(R3, key) do
                                t = zero(Float)
                                for p in -N:N
                                    (A == B && p == 0) && continue
                                    for q in -N:N
                                        ((A == C && q == 0) || (B == C && q == p)) && continue
                                        t += besselj(Int32(b-aâ€²-p), f*i_j) * besselj(Int32(c-b+p-q), f*k_l) * besselj(Int32(a-c+q), f*m_n) * (
                                            1 / 2(Îµâ‚€[Î±â€²] - Îµâ‚€[Î³] - q*Ï‰)     * (1/(Îµâ‚€[Î³] - Îµâ‚€[Î²]  - (p-q)*Ï‰) - 1/(Îµâ‚€[Î²] - Îµâ‚€[Î±â€²] + p*Ï‰)) +
                                            1 / 2(Îµâ‚€[Î±]  - Îµâ‚€[Î²] - p*Ï‰)     * (1/(Îµâ‚€[Î±] - Îµâ‚€[Î³]  - q*Ï‰)     - 1/(Îµâ‚€[Î³] - Îµâ‚€[Î²]  - (p-q)*Ï‰)) +
                                            1 / 6(Îµâ‚€[Î³]  - Îµâ‚€[Î²] - (p-q)*Ï‰) * (1/(Îµâ‚€[Î²] - Îµâ‚€[Î±â€²] + p*Ï‰)     + 1/(Îµâ‚€[Î±] - Îµâ‚€[Î³]  - q*Ï‰)) -
                                            1 / 3(Îµâ‚€[Î±]  - Îµâ‚€[Î³] - q*Ï‰) / (Îµâ‚€[Î²] - Îµâ‚€[Î±â€²] + p*Ï‰) )
                                    end
                                end
                                t
                            end
                            val *= s
                            H[Î±â€², Î±] += val
                        end
                    end
                end
            end
        end
    end
end

function push_state!(H_rows, H_cols, H_vals, val; row, col)
    push!(H_vals, val)
    push!(H_cols, col)
    push!(H_rows, row)
end

"Return key from the `R` dictionary; required for 2nd order (E)DPT."
function get_R!(R, U, Ï‰, f, nÎ±, d, i_j, k_l, aâ€², a, b, skipzero)
    key = (nÎ±, d, i_j, k_l, aâ€², a, b, skipzero)
    N = 20
    get!(R, key) do 
        s = zero(U)
        for n in -N:N
            skipzero && n == 0 && continue
            s += 1/(U*nÎ± - (d+n)*Ï‰) * besselj(Int32(-(aâ€²-b+n)), f*i_j) * besselj(Int32(a-b+n), f*k_l)
        end
        s
    end
end

"Return key from the `R` dictionary; required for 3rd order (E)DPT."
function get_R2!(R, U, Ï‰, f, Î”E1, Î”E2, d1, d2, J_indices, J_args, skipzero)
    i1, i2, i3 = J_indices
    x1, x2, x3 = J_args
    key = (Î”E1, Î”E2, d1, d2, i1, i2, i3, x1, x2, x3, skipzero)
    N = 20
    get!(R, key) do
        s = zero(U)
        for p in -N:N
            skipzero && p == 0 && continue
            s += 1 / (U*Î”E1 - (d1-p)*Ï‰) / (U*Î”E2 - (d2+p)*Ï‰) * 
                 besselj(Int32(i1), f*x1) * besselj(Int32(i2)-p, f*x2) * besselj(Int32(i3)+p, f*x3)
        end
        s
    end
end

"""
Calculate and return the spectrum for the values of ğ�‘ˆ in `Us`, using degenerate theory.
Calculation is based on the parameters in `bh`, including `bh.order`.
`bh.type` will be checked and set to `:edpt` if not set already.

If `sort=true`, the second returned argument, which is the permutation matrix, will be populated.
This allows one to isolate the quasienergies of states having the largest overlap with the ground state.
"""
function edpt(bh0::BoseHamiltonian{Float}, Us::AbstractVector{<:Real}; sort::Bool=false, showprogress=true) where {Float<:AbstractFloat}
    bh0.type != :edpt && update_params!(bh0; type=:edpt)
    
    nthreads = Threads.nthreads()
    if nthreads > 1
        nblas = BLAS.get_num_threads() # save original number of threads to restore later
        BLAS.set_num_threads(1)
    end
        
    nstates = size(bh0.H, 1)
    
    nU = length(Us)
    spectrum = Matrix{Float}(undef, nstates, nU)
    sp = Matrix{Int}(undef, nstates, nU) # sorting matrix

    bh_chnl = Channel{BoseHamiltonian{Float}}(nthreads)
    worskpace_chnl = Channel{HermitianEigenWs{Float, Matrix{Float}, Float}}(nthreads)
    for _ in 1:nthreads
        put!(bh_chnl, copy(bh0))
        put!(worskpace_chnl, HermitianEigenWs(bh0.H, vecs=sort))
    end

    progbar = ProgressMeter.Progress(nU; enabled=showprogress)
    Threads.@threads for i in eachindex(Us)
        bh = take!(bh_chnl)
        worskpace = take!(worskpace_chnl)
        update_params!(bh; U=Us[i])
        if sort
            spectrum[:, i], v = LAPACK.syevr!(worskpace, 'V', 'A', 'U', bh.H, 0.0, 0.0, 0, 0, 1e-6)
            sortperm!(@view(sp[:, i]), @view(v[1, :]), rev=true, by=abs2)
        else
            spectrum[:, i] = LAPACK.syevr!(worskpace, 'N', 'A', 'U', bh.H, 0.0, 0.0, 0, 0, 1e-6)[1]
        end
        put!(bh_chnl, bh)
        put!(worskpace_chnl, worskpace)
        ProgressMeter.next!(progbar)
    end
    ProgressMeter.finish!(progbar)
    nthreads > 1 && BLAS.set_num_threads(nblas) # restore original number of threads
    return spectrum, sp
end

"""
Calculate and return the spectrum for the values of ğ�‘ˆ in `Us`, using DPT.
`subspace` must contain the subspace number (as in `bh.space_of_state[:][1]`) of interest.
Calculation is based on the parameters in `bh`, including `bh.order`.
`bh.type` will be checked and set to `:dpt` if not set already.

If `sort=true`, the second returned argument, which is the permutation matrix, will be populated.
This allows one to isolate the quasienergies of states having the largest overlap with the ground state.
"""
function dpt(bh0::BoseHamiltonian{Float}, Us::AbstractVector{<:Real}; sort::Bool=false, showprogress=true, subspace::Integer=0) where {Float<:AbstractFloat}
    bh0.type != :dpt && update_params!(bh0; type=:dpt)

    nthreads = Threads.nthreads()
    if nthreads > 1
        nblas = BLAS.get_num_threads() # save original number of threads to restore later
        BLAS.set_num_threads(1)
    end
        
    As = findall(s -> s[1] == subspace, bh0.space_of_state) # `As` stores numbers of state that belong to space `subspace`
    nstates = length(As)
    nU = length(Us)
    spectrum = Matrix{Float}(undef, nstates, nU)
    sp = Matrix{Int}(undef, nstates, nU)

    bh_chnl = Channel{BoseHamiltonian{Float}}(nthreads)
    worskpace_chnl = Channel{HermitianEigenWs{Float, Matrix{Float}, Float}}(nthreads)
    H_chnl = Channel{Matrix{Float}}(nthreads)
    for _ in 1:nthreads
        put!(bh_chnl, copy(bh0))
        put!(worskpace_chnl, HermitianEigenWs(Matrix{Float}(undef, nstates, nstates), vecs=sort))
        put!(H_chnl, Matrix{Float}(undef, nstates, nstates))
    end

    progbar = ProgressMeter.Progress(nU; enabled=showprogress)
    Threads.@threads for i in eachindex(Us)
        nanfound = false
        bh = take!(bh_chnl)
        worskpace = take!(worskpace_chnl)
        H = take!(H_chnl)
        
        update_params!(bh; U=Us[i])

        for i in eachindex(As), j in i:length(As)
            h = bh.H[As[j], As[i]]
            isnan(h) && (nanfound = true; break)
            H[j, i] = bh.H[As[j], As[i]]
        end

        # if `Us[i]` is such that DPT is invalid, then `bh.H` is (or is close to being) singular, so that Inf's will appear during diagonalisation.
        if nanfound
            spectrum[:, i] .= NaN
        else
            if sort
                spectrum[:, i], v = LAPACK.syevr!(worskpace, 'V', 'A', 'L', H, 0.0, 0.0, 0, 0, 1e-6)
                !isnan(spectrum[1, i]) && sortperm!(@view(sp[:, i]), @view(v[1, :]), rev=true, by=abs2)
            else
                spectrum[:, i] = LAPACK.syevr!(worskpace, 'N', 'A', 'L', H, 0.0, 0.0, 0, 0, 1e-6)[1]
            end
        end
        put!(bh_chnl, bh)
        put!(worskpace_chnl, worskpace)
        put!(H_chnl, H)
        ProgressMeter.next!(progbar)
    end
    ProgressMeter.finish!(progbar)
    nthreads > 1 && BLAS.set_num_threads(nblas) # restore original number of threads
    return spectrum, sp
end

"""
Analyse residual couplings with states outside the FZ.
The passed `bh` has to be initialised with the required Ï‰â‚— defining the FZ.
Take a state |ğ�›¼0âŸ©âŸ© together with the states |ğ�›¼â€²ğ�‘›âŸ©âŸ© which it is coupled to in first order. Scan ğ�‘› to find the strongest coupling ratio defined as
<coupling strength> / <energy distance>. Save this ratio to the correponding elements `bh.H[Î±â€², Î±]`. Repeat for all (Î±, Î±â€²) pairs.
Additionally, return a matrix `W` with ğ�‘›'s: `W[Î±â€², Î±] = n`.
"""
function residuals!(bh::BoseHamiltonian{Float}) where {Float<:AbstractFloat}
    (;index_of_state, ncells, neis_of_cell) = bh.lattice
    (;J, f, Ï‰, Îµâ‚€, space_of_state, H) = bh

    H .= 0
    W = zeros(Int, size(H))

    bra = similar(bh.lattice.basis_states[1])
    # take each basis state and find which transitions are possible
    for (ket, Î±) in index_of_state
        _, a = space_of_state[Î±]
        for i = 1:ncells # iterate over the terms of the Hamiltonian
            for (j, i_j) in neis_of_cell[i]
                if (ket[j] > 0) # check that a particle is present at site `j` so that destruction ğ�‘�â±¼ is possible
                    copy!(bra, ket)
                    bra[j] -= 1
                    bra[i] += 1
                    Î±â€² = index_of_state[bra]
                    _, aâ€² = space_of_state[Î±â€²]
                    r_max, n_max = 0.0, 0
                    for n in -5:5 # large values of `n` are likely to lead to low ratios because of large energy distance
                        n == 0 && continue # skip levels inside the FZ
                        r = besselj(a - (aâ€² + n), f*i_j) / (Îµâ‚€[Î±] - (Îµâ‚€[Î±â€²] - n*Ï‰)) |> abs # `n`s are with different signs because adding `n` to subspace number means subtracting `nÏ‰` from the energy
                        r > r_max && (r_max = r; n_max = -n) # -n will show the number of Floquet zone
                    end
                    H[Î±â€², Î±] = J * r_max * sqrt( (ket[i]+1) * ket[j] )
                    W[Î±â€², Î±] = n_max
                end
            end
        end
    end
    return W
end

function quasienergy(bh::BoseHamiltonian{Float}, Us::AbstractVector{<:Real}; outdir::String="", sort::Bool=false, showprogress::Bool=true, reltol::Real=1e-3) where {Float<:AbstractFloat}
    nstates = length(bh.lattice.basis_states)
    nU = length(Us)
    if nprocs() == 1
        Îµ = Matrix{Float}(undef, nstates, nU)
        sp = Matrix{Int}(undef, nstates, nU)
    else
        Îµ = SharedMatrix{Float}(nstates, nU)
        sp = SharedMatrix{Int}(nstates, nU)
    end
    quasienergy!(Îµ, sp, bh, Us; outdir, sort, showprogress, reltol)
    return Îµ, sp
end

"""
Calculate quasienergy spectrum of `bh` via monodromy matrix for each value of ğ�‘ˆ in `Us`.
Loop over `Us` is parallelised: using processes if Julia is launched with more than one process, and using threads otherwise.
BLAS threading is turned off, but is restored to the original state upon finishing the calculation.

If `outdir` is passed, a new directory will be created (if it does not exist) where the quasienergy spectrum at each `i`th value of `Us`
will be output immediately after calculation. The files are named as "<i>.txt"; the first value in the file is `Us[i]`, and the following
are the quasienergies.

If `sort=true`, the second argument (`sp`), which is the permutation matrix, will be populated.
This allows one to isolate the quasienergies of states having the largest overlap with the ground state.
If output is enabled using `outdir`, the files with have a second column where the first value is `Us[i]`, and the following are the permutation integers.

`reltol` is passed to the ODE solver.
"""
function quasienergy!(Îµ::AbstractMatrix, sp::AbstractMatrix, bh::BoseHamiltonian{Float}, Us::AbstractVector{<:Real}; outdir::String="", sort::Bool=false,
                      showprogress::Bool=true, reltol::Real=1e-3) where {Float<:AbstractFloat}
    (;J, f, Ï‰, Eâ‚€) = bh
    Cmplx = (Float == Float32 ? ComplexF32 : ComplexF64)
    (;index_of_state, ncells, neis_of_cell) = bh.lattice
    
    H_rows, H_cols, H_vals = Int[], Int[], Cmplx[]
    H_sign = Float[] # stores the sign of the tunneling phase for each off-diagonal element

    # Fill the off-diagonal elemnts of the Hamiltonian for `f` = 0
    bra = similar(bh.lattice.basis_states[1])
    for (ket, index) in index_of_state
        for i = 1:ncells # iterate over the terms of the Hamiltonian
            for (j, s) in neis_of_cell[i]
                # ğ�‘�â€ áµ¢ ğ�‘�â±¼
                if (ket[j] > 0) # check that a particle is present at site `j` so that destruction ğ�‘�â±¼ is possible
                    val = -im * -J * sqrt( (ket[i]+1) * ket[j] ) # multiply by `-im` as on the rhs of âˆ‚â‚œğ�œ“ = -iğ��»ğ�œ“
                    copy!(bra, ket)
                    bra[j] -= 1
                    bra[i] += 1
                    row = index_of_state[bra]
                    push_state!(H_rows, H_cols, H_vals, val; row, col=index)
                    push!(H_sign, s)
                end
            end
        end
    end

    nstates = size(bh.H, 1)

    # append placeholders for storing diagonal elements
    append!(H_vals, zeros(Float, nstates))
    append!(H_sign, zeros(Float, nstates))
    append!(H_rows, 1:nstates)
    append!(H_cols, 1:nstates)

    nthreads = Threads.nthreads()
    np = nprocs()
    noblas = nthreads > 1 || np > 1
    if noblas
        nblas = BLAS.get_num_threads() # save original number of threads to restore later
        BLAS.set_num_threads(1)
    end

    Câ‚€ = Matrix{Cmplx}(I, nstates, nstates)
    
    T = 2Ï€ / Ï‰
    tspan = (0.0, T)
    
    S = sparse(H_rows, H_cols, H_sign) # construct the "sign" matrix
    H_sign_vals = nonzeros(S) # get the values so that they are in the same order as `H_base_vals`

    # construct the base matrix, which will not be mutated
    H_base = sparse(H_rows, H_cols, H_vals)
    H_rows, H_cols, H_base_vals = findnz(H_base)
    dind = (H_rows .== H_cols)

    outdir != "" && mkpath(outdir)

    params = (H_base, H_base_vals, H_base_vals, H_sign_vals, Ï‰, f) # the contents of `params` do not matter as they will be replaced
    prob = ODEProblem(schrodinger!, Câ‚€, tspan, params, save_everystep=false, save_start=false)

    if np > 1 # distributed mode
        nw = nworkers()
        pmap(1:nw) do pid
            # diagonal of `H_buff` will remain equal to the diagonal of `H_base` throughout diffeq solving, while off-diagnoal elements will be mutated at each step
            H_buff = H_base # `H_base` will be copied to each process, so we only set `H_buff` to reference it
            integrator = OrdinaryDiffEq.init(prob, Tsit5(); reltol)
            workspace = EigenWs(Câ‚€, rvecs=sort)
            for i in pid:nw:length(Us)
                quasienergy_step!(Îµ, sp, i, Us[i], dind, Eâ‚€, H_buff, integrator, workspace, Câ‚€, H_base_vals, H_sign_vals, Ï‰, f, sort, outdir)
            end
        end
    else # threaded mode
        H_buff_chnl = Channel{typeof(H_base)}(nthreads)
        worskpace_chnl = Channel{EigenWs{Cmplx, Matrix{Cmplx}, Float}}(nthreads)
        integrator_chnl = Channel{typeof(OrdinaryDiffEq.init(prob, Tsit5(); reltol))}(nthreads)
        for _ in 1:nthreads
            put!(H_buff_chnl, copy(H_base))
            put!(worskpace_chnl, EigenWs(Câ‚€, rvecs=sort))
            put!(integrator_chnl, OrdinaryDiffEq.init(prob, Tsit5(); reltol))
        end

        progbar = ProgressMeter.Progress(length(Us); enabled=showprogress)
        Threads.@threads for i in eachindex(Us)
            H_buff = take!(H_buff_chnl)
            integrator = take!(integrator_chnl)
            workspace = take!(worskpace_chnl)
            
            quasienergy_step!(Îµ, sp, i, Us[i], dind, Eâ‚€, H_buff, integrator, workspace, Câ‚€, H_base_vals, H_sign_vals, Ï‰, f, sort, outdir)

            put!(H_buff_chnl, H_buff)
            put!(integrator_chnl, integrator)
            put!(worskpace_chnl, workspace)

            ProgressMeter.next!(progbar)
        end
        ProgressMeter.finish!(progbar)
    end
    noblas && BLAS.set_num_threads(nblas) # restore original number of threads
end

function quasienergy_step!(Îµ, sp, i, U, dind, Eâ‚€, H_buff, integrator, workspace, Câ‚€, H_base_vals, H_sign_vals, Ï‰, f, sort, outdir)
    # diagonal of `H_buff` will remain equal to the diagonal of `H_base` throughout diffeq solving, while off-diagnoal elements will be mutated at each step
    H_buff_vals = nonzeros(H_buff) # a view to non-zero elements
    H_buff_vals[dind] .= U .* (-im .* Eâ‚€) # update diagonal of the Hamiltonian
    
    reinit!(integrator, Câ‚€)
    integrator.p = (H_buff, H_buff_vals, H_base_vals, H_sign_vals, Ï‰, f)
    sol = solve!(integrator)

    if sort
        t = LAPACK.geevx!(workspace, 'N', 'N', 'V', 'N', sol[end])
        e, v = t[2], t[4] # convenience views
        sortperm!(@view(sp[:, i]), @view(v[1, :]), rev=true, by=abs2)
        @. Îµ[:, i] = -Ï‰ * angle(e) / 2Ï€
    else
        LAPACK.geevx!(workspace, 'N', 'N', 'N', 'N', sol[end])
        @. Îµ[:, i] = -Ï‰ * angle(workspace.W) / 2Ï€
    end

    if outdir != ""
        open(joinpath(outdir, "$(i).txt"), "w") do io
            if sort
                writedlm(io, vcat([U U], [Îµ[:, i] sp[:, i]]))
            else
                writedlm(io, vcat([U], Îµ[:, i]))
            end
        end
    end
end

"SchrÃ¶dinger equation used for monodromy matrix calculation."
function schrodinger!(du, u, params, t)
    H_buff, H_buff_vals, H_base_vals, H_sign_vals, Ï‰, f = params
    # update off diagonal elements of the Hamiltonian
    p = cis(f * sin(Ï‰*t)); n = cis(-f * sin(Ï‰*t));
    for (i, s) in enumerate(H_sign_vals)
        if s > 0 
            H_buff_vals[i] = H_base_vals[i] * p
        elseif s < 0
            H_buff_vals[i] = H_base_vals[i] * n
        end
    end
    mul!(du, H_buff, u)
end

"""
Calculate quasienergy spectrum via monodromy matrix for each value of ğ�‘ˆ in `Us`.
`bh` is used as a parameter holder, but `bh.U`, `bh.type`, and `bh.order` do not matter.
By default, loop over `Us` is parallelised using all threads that julia was launched with: `nthreads=Threads.nthreads()`.
Setting `nthreads=1` makes the loop over `Us` sequential, but the diffeq solving uses BLAS threading.
For `nthreads > 1`, BLAS threading is turned off, but is restored to the original state upon finishing the calculation.
"""
function quasienergy_dense(bh::BoseHamiltonian{Float}, Us::AbstractVector{<:Real}; nthreads::Int=Threads.nthreads()) where {Float<:AbstractFloat}
    (;J, f, Ï‰, Eâ‚€) = bh
    (;index_of_state, ncells, neis_of_cell) = bh.lattice
    Cmplx = (Float == Float32 ? ComplexF32 : ComplexF64)

    nstates = size(bh.H, 1)
    H = zeros(Cmplx, nstates, nstates)
    H_sign = zeros(Int, nstates, nstates)

    # Construct the Hamiltonian with `f` = 0
    # off-diagonal elements ğ�‘�â€ áµ¢ ğ�‘�â±¼
    for (ket, Î±) in index_of_state
        for i = 1:ncells # iterate over the terms of the Hamiltonian
            for (j, s) in neis_of_cell[i]
                # ğ�‘�â€ áµ¢ ğ�‘�â±¼
                if (ket[j] > 0) # check that a particle is present at site `j` so that destruction ğ�‘�â±¼ is possible
                    val = -im * -J * sqrt( (ket[i]+1) * ket[j] ) # multiply by `-im` as on the rhs of âˆ‚â‚œğ�œ“ = -iğ��»ğ�œ“
                    bra = copy(ket)
                    bra[j] -= 1
                    bra[i] += 1
                    Î±â€² = index_of_state[bra]
                    H[Î±â€², Î±] = val
                    H_sign[Î±â€², Î±] = s
                end
            end
        end
    end

    n_U = length(Us)
    Îµ = Matrix{Float}(undef, nstates, n_U)
    Câ‚€ = Matrix{Cmplx}(I, nstates, nstates)
    
    T = 2Ï€ / Ï‰
    tspan = (0.0, T)

    if nthreads > 1
        n_blas = BLAS.get_num_threads() # save original number of threads to restore later
        BLAS.set_num_threads(1)
    end
    executor = FLoops.ThreadedEx(basesize=length(Us)Ã·nthreads)

    # progbar = ProgressMeter.Progress(length(Us))  # slows down conputation up to 1.5 times!

    @floop executor for (i, U) in enumerate(Us)
        @init begin
            # diagonal of `H_buff` will remain equal to -ğ�‘–ğ�‘ˆ times the diagonal of `H` throughout diffeq solving,
            # while off-diagnoal elemnts will be mutated at each step
            H_buff = zeros(Cmplx, nstates, nstates)
        end
        H_buff[diagind(H_buff)] .= U .* (-im .* Eâ‚€)
        params = (H_buff, H, H_sign, Ï‰, f)
        prob = ODEProblem(schrodinger_dense!, Câ‚€, tspan, params, save_everystep=false)
        sol = solve(prob, Tsit5())
        Îµ[:, i] = -Ï‰ .* angle.(eigvals(sol[end])) ./ 2Ï€

        # ProgressMeter.next!(progbar)
    end
    # ProgressMeter.finish!(progbar)
    nthreads > 1 && BLAS.set_num_threads(n_blas) # restore original number of threads

    return Îµ
end

"Dense version of the SchrÃ¶dinger equation used for monodromy matrix calculation."
function schrodinger_dense!(du, u, params, t)
    H_buff, H_base, H_sign, Ï‰, f = params
    p = cis(f * sin(Ï‰*t)); n = cis(-f * sin(Ï‰*t));
    for (i, s) in enumerate(H_sign)
        if s > 0 
            H_buff[i] = H_base[i] * p
        elseif s < 0
            H_buff[i] = H_base[i] * n
        end
    end
    mul!(du, H_buff, u)
end