text:
include("../models/StochasticProcess.jl")
include("../models/AssetModel.jl")

using LinearAlgebra, Interpolations

struct HybridModel{T<:AbstractFloat,D<:StochasticProcess,A<:AssetModel,F<:StochasticProcess} <: StochasticProcess
    domAlias::String                          #  name of our domestic (numeraire) currency
    domRatesModel::D                          #  domestic rates model specifies numeraire
    forAliases::Array{String}                 #  list of foreign currencies (all relative to dom currency)
    forAssetModels::Array{A}                  #  list of foreign asset models
    forRatesModels::Array{F}                  #  list of foreign rates models
    correlations::Array{T,2}                  #  2d array of instantanous correlations
    # additional data
    index::Dict{String,Int64}                 # a dictionary from foreign alias to index
    _size::Int64                              # sum of model states
    _factors::Int64                           # sum of model _factors
    modelsStartIdx::Array{Int64}              # were do foreign models start in X
    L::LowerTriangular{Float64,Array{Float64,2}}   # L L^T = correlations
    hybAdjTimes::Array{T}                    # array of reference times for adjuster
    hybVolAdj::Array{T}                      # array of the additive volatility adjuster
end

# constructor

function HybridModel(domAlias, domRatesModel,
    forAliases, forAssetModels, forRatesModels, correlations)
    # add sanity checks here
    # we need to know the model index for a given alias
    index = Dict([ (forAliases[k],k) for k = 1:size(forAliases)[1] ])
    # manage model indices in state variable
    _size = stateSize(domRatesModel)
    _factors = factors(domRatesModel)
    modelsStartIdx = Vector{Int64}()
    lastModelIdx = stateSize(domRatesModel)
    for (assetModel, ratesModel) in zip(forAssetModels,forRatesModels)
        _size += (stateSize(assetModel) + stateSize(ratesModel))
        _factors += (factors(assetModel) + factors(ratesModel))
        push!(modelsStartIdx,lastModelIdx)
        push!(modelsStartIdx,lastModelIdx + stateSize(assetModel))
        lastModelIdx += (stateSize(assetModel) + stateSize(ratesModel))
    end
    if  !isnothing(correlations)
        L = cholesky(correlations).L
    else
        L = nothing
    end
    hybAdjTimes = zeros(0)
    hybVolAdj = zeros(0)
    return HybridModel(domAlias,domRatesModel,forAliases,forAssetModels,forRatesModels,
        correlations,index,_size,_factors,modelsStartIdx,L,hybAdjTimes,hybVolAdj)
end

HybridModel(model::HybridModel,hybAdjTimes,hybVolAdj) =
    HybridModel(model.domAlias, model.domRatesModel, model.forAliases, model.forAssetModels,
        model.forRatesModels, model.correlations, model.index, model._size, model._factors,
        model.modelsStartIdx, model.L, hybAdjTimes, hybVolAdj)

# stochastic process interface

function stateSize(self::HybridModel)   # dimension of X(t)
    return self._size
end

function factors(self::HybridModel)   # dimension of W(t)
    return self._factors
end

function initialValues(self::HybridModel)
    X0 = initialValues(self.domRatesModel)
    for (assetModel, ratesModel) in zip(self.forAssetModels,self.forRatesModels)
        X0 = append!(X0, initialValues(assetModel))
        X0 = append!(X0, initialValues(ratesModel))
    end
    return X0
end

# evolve X(t0) -> X(t0+dt) using independent Brownian increments dW
# t0, dt are assumed float, X0, X1, dW are np.array
@views function evolve(self::HybridModel, t0, X0, dt, dW, X1)
    if !isnothing(self.L)
        dZ = self.L * dW
    else
        dZ = dW
    end
    if isnothing(self.forAliases)  # take a short cut
        evolve(self.domRatesModel,t0, X0, dt, dZ, X1)
        return nothing
    end
    # evolve domestic rates
    domSize = stateSize(self.domRatesModel)  # shorten expressions
    evolve(self.domRatesModel, t0, X0[begin:domSize], dt, dZ[begin:factors(self.domRatesModel)], X1[begin:domSize])
    # we need the domestic drift r_d
    r_d = shortRateOverPeriod(self.domRatesModel, t0, dt, X0[begin:domSize], X1[begin:domSize])
    # now we iterate over foreign models
    corrStartIdx = factors(self.domRatesModel)  #  we need to keep track of our sub-correlations
    for (k, alias) in enumerate(self.forAliases)
        # carefully collect Brownian increments
        startIdx = corrStartIdx + 1
        endIdx   = startIdx + factors(self.forAssetModels[k]) - 1
        dw_asset = dZ[ startIdx : endIdx ]
        #
        startIdx = endIdx + 1
        endIdx   = startIdx + factors(self.forRatesModels[k]) - 1
        dw_rates = dZ[ startIdx : endIdx ]
        # we need the starting point states for evolution, y0 (asset), x0 (rates)
        startIdx = self.modelsStartIdx[2 * k - 1] + 1
        endIdx   = startIdx + stateSize(self.forAssetModels[k]) - 1
        y0 = X0[ startIdx : endIdx ]
        y1 = X1[ startIdx : endIdx ]
        #
        startIdx = endIdx + 1
        endIdx   = startIdx + stateSize(self.forRatesModels[k]) - 1
        x0 = X0[ startIdx : endIdx ]
        x1 = X1[ startIdx : endIdx ]
        #
        # Quanto adjustment
        # we use the model-independent implementation to allow for credit hybrid components
        # todo: capture case self.correlations = None
        colIdx   = corrStartIdx + 1
        startIdx = colIdx   + factors(self.forAssetModels[k])
        endIdx   = startIdx + factors(self.forRatesModels[k]) - 1
        qAdj = self.correlations[colIdx, startIdx : endIdx]
        # we want to modify qAdj vector w/o changing the correlation
        qAdj_ = Array(qAdj)
        # we need to extend the input state for our asset mode to account for drift and adjuster
        y0_ = append!(Array(y0), [0.0, 0.0])
        y0_[end] = hybridVolAdjuster(self, k, t0)  # implement adjuster!
        assetVol = volatility(self.forAssetModels[k], t0, y0_)
        qAdj_ *= (assetVol*sqrt(dt))
        dw_rates_ = dw_rates - qAdj_  #  create a new vector
        # evolve foreign rates
        evolve(self.forRatesModels[k], t0, x0, dt, dw_rates_, x1)
        # calculate FX drift, volAdjuster and extend input state
        r_f = shortRateOverPeriod(self.forRatesModels[k], t0, dt, x0, x1)
        y0_[end-1] = r_d - r_f   # FX drift
        # finally we can evolve FX
        evolve(self.forAssetModels[k], t0, y0_, dt, dw_asset, y1)
        # no need to copy results coz we work on views of X1
        # but we need to update stochastic factor index
        corrStartIdx += (factors(self.forAssetModels[k]) + factors(self.forRatesModels[k]))
    end
    return nothing
end

# interface for payoff calculation

@views function numeraire(self::HybridModel, t, X)
    return numeraire(self.domRatesModel,t,X[begin:stateSize(self.domRatesModel)])
end

@views function asset(self::HybridModel, t, X, alias)
    if isnothing(alias)
        return 1.0
    end
    if alias==self.domAlias
        return 1.0
    end
    k = self.index[alias]  # this should throw an exception if alias is unknown
    startIdx = self.modelsStartIdx[2 * k - 1] + 1
    endIdx   = startIdx + stateSize(self.forAssetModels[k]) - 1
    y = X[ startIdx : endIdx ]
    return asset(self.forAssetModels[k], t, y, alias)
end

@views function zeroBond(self::HybridModel, t, T, X, alias)
    if isnothing(alias) || alias==self.domAlias
        x = X[begin:stateSize(self.domRatesModel)]
        return zeroBond(self.domRatesModel, t, T, x, alias)
    end
    k = self.index[alias]  # this should throw an exception if alias is unknown
    startIdx = self.modelsStartIdx[2 * k] + 1
    endIdx   = startIdx + stateSize(self.forRatesModels[k]) - 1
    x = X[ startIdx : endIdx ]
    return zeroBond(self.forRatesModels[k], t, T, x, alias)
end

# adjuster methodology for stochastic rates and FX volatility

function hybridVolAdjuster(self::HybridModel, forIdx, t)
    if size(self.hybAdjTimes)==(0,)
        return 0.0  # default
    end
    # linear interpolation with constant extrapolation
    # maybe better use scipy interpolation with linear extraplation
    interp = LinearInterpolation(self.hybAdjTimes,@view self.hybVolAdj[forIdx,:])
    return interp(t)
end

function hybridVolAdjuster(self::HybridModel, hybAdjTimes::Array)
    if !(size(hybAdjTimes)[1]>1)
        throw(ArgumentError("size(hybAdjTimes)[1]>1 required."))
    end
    if !(hybAdjTimes[1]==0.0)
        throw(ArgumentError("hybAdjTimes[0]==0.0 required."))
    end
    for k = 2:size(hybAdjTimes)[1]
        if !(hybAdjTimes[k]>hybAdjTimes[k-1])
            throw(ArgumentError("hybAdjTimes[k]>hybAdjTimes[k-1] required."))
        end
    end
    # initialise
    localVol = zeros((size(self.forAliases)[1],size(hybAdjTimes)[1]))
    hybrdVol = ones((size(self.forAliases)[1],size(hybAdjTimes)[1])) #  1.0 is required for p calculation
    hybVolAdj = zeros((size(self.forAliases)[1],size(hybAdjTimes)[1]))
    S0 = [ asset(m, 0.0, initialValues(m), nothing) for m in self.forAssetModels ]
    #
    # calculate vols at zero
    for i = 1:size(S0)[1]
        y0 = append!(initialValues(self.forAssetModels[i]), [0.0, 0.0])
        localVol[i,1] = volatility(self.forAssetModels[i], 0.0, y0)
        hybrdVol[i,1] = localVol[i,1]
    end
    # now we start with the actual methodology...
    corrStartIdx = factors(self.domRatesModel)
    for i = 1:size(S0)[1]
        # we collect all relevant correlations
        # recall,
        #   Y0 is domestic rates model
        #   X1 is asset (or FX) model
        #   Y1 is forign rates model
        #
        colIdx   = corrStartIdx + 1
        # domestic rates vs FX, ASSUME vol-FX correlation is zero
        rhoY0X1 = self.correlations[ begin:factors(self.domRatesModel), colIdx]   #
        # foreign rates vs FX, ASSUME vol-FX correlation is zero
        startIdx = colIdx   + factors(self.forAssetModels[i])
        endIdx   = startIdx + factors(self.forRatesModels[i]) - 1
        # rhoX1Y1 = self.correlations[colIdx, startIdx : endIdx]
        rhoY1X1 = self.correlations[startIdx : endIdx, colIdx]  # use col-vector to avoid transpose
        # rates vs rates, ASSUME all vol-... correlation are zero
        rhoY0Y1 = self.correlations[ begin:factors(self.domRatesModel), startIdx : endIdx ]
        # update stochastic factor index
        corrStartIdx += (factors(self.forAssetModels[i]) + factors(self.forRatesModels[i]))
        # bootstrap over adjuster times
        for (k, T) in [ e for e in enumerate(hybAdjTimes) ][2:end]
            # ATM forward and effective local volatility
            dfDom = zeroBond(self.domRatesModel, 0.0, T, initialValues(self.domRatesModel), nothing)
            dfFor = zeroBond(self.forRatesModels[i], 0.0, T, initialValues(self.forRatesModels[i]), nothing)
            S = S0[i] * dfFor / dfDom  # maybe it's worth to save S for debugging
            y = zeros(stateSize(self.forAssetModels[i]) + 2)  # this is asset model-dependent
            y[1] = log(S / S0[i])
            localVol[i,k] = volatility(self.forAssetModels[i],hybAdjTimes[k], y)
            # calculate derivative of hybrid variance
            hPrime = zeros(k)
            for (j, t) in enumerate(hybAdjTimes[begin:k])
                sigmaP0      = zeroBondVolatility(self.domRatesModel, t, T)
                sigmaP0Prime = zeroBondVolatilityPrime(self.domRatesModel, t, T)
                sigmaP1      = zeroBondVolatility(self.forRatesModels[i], t, T)
                sigmaP1Prime = zeroBondVolatilityPrime(self.forRatesModels[i], t, T)
                #
                sigma0 = sigmaP0 - rhoY0Y1*sigmaP1 + hybrdVol[i,j]*rhoY0X1  # bootstrapping enters here
                sum0 = transpose(sigmaP0Prime) * sigma0
                #
                sigma1 = sigmaP1 - transpose(rhoY0Y1)*sigmaP0 - hybrdVol[i,j]*rhoY1X1  # bootstrapping enters here
                sum1 = transpose(sigma1) * sigmaP1Prime
                # collect terms and finish
                hPrime[j] = 2.0*(sum0 + sum1)  # this will likely fail for deterministic models
            end
            p = 0.5 * hPrime[k] * (hybAdjTimes[k] - hybAdjTimes[k - 1])
            q = 0.5 * hPrime[k-1] * (hybAdjTimes[k] - hybAdjTimes[k - 1])
            for j = 2:k-1
                q += 0.5 * (hPrime[j - 1] + hPrime[j]) * (hybAdjTimes[j] - hybAdjTimes[j - 1])
            end
            # let's see if this works...
            root2 = p*p / 4.0 - q + localVol[i,k] * localVol[i,k]
            if !(root2>=0.0)
                throw(ArgumentError("root2>=0.0 required."))
            end
            hybrdVol[i,k] = -p / 2.0 + sqrt(root2)
            if !(hybrdVol[i,k]>0.0)
                throw(ArgumentError("hybrdVol[i,k]>0.0 required."))
            end
            # maybe we should add some more safety checks here...
            hybVolAdj[i,k] = hybrdVol[i,k] - localVol[i,k]
        end
    end
    #
    return hybVolAdj
end

max_stars_repo_name: sschlenkrich/HybridMonteCarlo

max_issues_repo_name: sschlenkrich/HybridMonteCarlo

max_forks_repo_name: sschlenkrich/HybridMonteCarlo


text:
module GeometryBasics

using StaticArrays, Tables, StructArrays

using Base: @propagate_inbounds

include("fixed_arrays.jl")
include("basic_types.jl")
include("metadata.jl")
include("viewtypes.jl")

end # module

ext: jl


text:
# Range operations
# Take two runends vectors (strictly increasing uints) and find the number of unique values for the disjoin operation
function disjoin_length(x::Vector, y::Vector)
    i = length(x)
    j = length(y)
    nrun = i + j
    @inbounds while i > 0 && j > 0
        if x[i] > y[j]
            i = i - 1
        elseif x[i] < y[j]
            j = j - 1
        else
            i = i - 1
            j = j - 1
            nrun = nrun - 1
        end
    end
    return(nrun)
end
disjoin_length(x::RLEVector, y::RLEVector) = disjoin_length(x.runends, y.runends)

"""
Takes runends from two RLEVectors, make one new runends breaking the pair into non-overlapping runs.
Basically, this is an optimized `sort!(unique([x,y])))`. This is useful when comparing two RLEVector
objects. The values corresponding to each disjoint run in `x` and `y` can then be compared directly.

## Returns
An integer vector, of a type that is the promotion of the eltypes of the runends of x and y.

    ## Examples
    x = RLEVector([1,1,2,2,3,3])
    y = RLEVector([1,1,1,2,3,4])
    for (i,j) in disjoin(x,y)
        println(x[i] + y[j])
    end
    """
function disjoin(x::Vector,  y::Vector)
    length(x) == 0 && return(y) # At least one value to work on
    length(y) == 0 && return(x) # At least one value to work on
    nrun = disjoin_length(x, y)
    i = length(x)
    j = length(y)
    runends = Array(promote_type(eltype(x), eltype(y)), nrun)
    @inbounds while true
        xi = x[i]
        yj = y[j]
        if xi > yj
            runends[nrun] = xi
            i = i - 1
        elseif xi < yj
            runends[nrun] = yj
            j = j - 1

        else
            runends[nrun] = xi
            i = i - 1
            j = j - 1
        end
        nrun = nrun - 1
        if i == 0
            for r in 1:j runends[r] = y[r] end
            break
        elseif j == 0
            for r in 1:i runends[r] = x[r] end
            break
        end
    end
    return(runends)
end

function disjoin(x::RLEVector, y::RLEVector)
    length(x) != length(y) && error("RLEVectors of unequal length.")
    runends = disjoin(x.runends, y.runends)
    runvalues_x = x.runvalues[ ind2run(x, runends) ]
    runvalues_y = y.runvalues[ ind2run(y, runends) ]
#    #nrun = disjoin_length(x.runends, y.runends) # Alternative
#    #nx = nrun(x)
#    #ny = nrun(y)
#    xv = x.runvalues
#    yv = y.runvalues
#    xe = x.runends
#    ye = y.runends
#    runends = disjoin(xe, ye)
#    nrun = length(runends)
#
#    runvalues_x = Array(eltype(x), nrun)
#    runvalues_y = Array(eltype(x), nrun)
#    runvalues_x[:] = 0 # debug
#    runvalues_y[:] = 0 # debug
#    i = j = 1
#    while true
#        if xe[i] > ye[j]
#            println("i bigger")
#            current_end = runends[runind] = xe[i]
##            runvalues_x[runind] = xv[i]
##            runvalues_y[runind] = yv[j]
##            i = i - 1
#        elseif xe[i] < ye[j]
#            println("j bigger")
#            runends[runind] = ye[j]
##            runvalues_x[runind] = xv[i]
##            runvalues_y[runind] = yv[j]
##            j = j - 1
#        else
#            println("tie")
#            runends[runind] = xe[i]
##            runvalues_x[runind] = xv[i]
##            runvalues_y[runind] = yv[j]
##            i = i - 1
##            j = j - 1
#        end
#
#        runind = runind - 1
#        println(runvalues_x)
#        println(runvalues_y)
#    end
    return( (runends, runvalues_x, runvalues_y) )
end

# optimization opportunities: hoist rle element lookups and use the searchsortedfirst with all the args
function rangeMeans(ranges::Vector{UnitRange}, rle::RLEVector)
    res = similar(ranges, Float64)
    first_run = 1
    last_run = nrun(rle)
    @inbounds for (i, r) in enumerate(ranges)
        first_ind = r[1]
        last_ind = r[end]
        first_run = searchsortedfirst(rle.runends, first_ind)
        last_run = searchsortedlast(rle.runends, last_ind)
        if first_run == last_run  # Range all in one run, special case here allows simpler logic below
            res[i] = rle.runvalues[first_run]
        else
            # First run
            current_end = rle.runends[first_run]
            inner_n = (current_end - first_ind) + 1
            temp_sum = rle.runvalues[first_run] * inner_n
            # Inner runs
            for run_index in (first_run + 1):(last_run - 1)
                previous_end = current_end
                current_end = rle.runends[run_index]
      	        inner_n = current_end - previous_end
	        temp_sum = temp_sum + (rle.runvalues[run_index] * inner_n)
            end
            # Last run
            inner_n = last_ind - current_end
            temp_sum = temp_sum + (rle.runvalues[last_run] * inner_n)
            # Calculate mean
            res[i] = temp_sum / ((last_ind - first_ind) + 1)
        end
    end
    return(res)
end


text:
# NOTE: ArrayReg shares some common interfaces with Array

using LinearAlgebra

export isnormalized, normalize!, regadd!, regsub!, regscale!, norm

"""
    isnormalized(r::ArrayReg) -> Bool

Check if the register is normalized.
"""
isnormalized(r::ArrayReg) =
    all(sum(copy(r) |> relax!(to_nactive = nqubits(r)) |> probs, dims = 1) .≈ 1)
isnormalized(r::AdjointArrayReg) = isnormalized(parent(r))

"""
    normalize!(r::ArrayReg)

Normalize the register `r` in-place by its `2`-norm.
"""
function LinearAlgebra.normalize!(r::ArrayReg{B}) where {B}
    batch_normalize!(reshape(r.state, :, B))
    return r
end

LinearAlgebra.normalize!(r::AdjointArrayReg) = (normalize!(parent(r)); r)

LinearAlgebra.norm(r::ArrayReg{1}) = norm(statevec(r))
LinearAlgebra.norm(r::ArrayReg{B}) where B = [norm(view(reshape(r.state, :, B), :, ib)) for ib=1:B]

# basic arithmatics

# neg
Base.:-(reg::ArrayReg) = ArrayReg(-state(reg))
Base.:-(reg::AdjointArrayReg) = adjoint(-parent(reg))

# +, -
for op in [:+, :-]
    @eval function Base.$op(lhs::ArrayReg{B}, rhs::ArrayReg{B}) where {B}
        return ArrayReg(($op)(state(lhs), state(rhs)))
    end

    @eval function Base.$op(
        lhs::ArrayReg{B,T1,<:Transpose},
        rhs::ArrayReg{B,T2,<:Transpose},
    ) where {B,T1,T2}
        return ArrayReg(transpose(($op)(state(lhs).parent, state(rhs).parent)))
    end

    @eval function Base.$op(lhs::AdjointArrayReg{B}, rhs::AdjointArrayReg{B}) where {B}
        r = $op(parent(lhs), parent(rhs))
        return adjoint(r)
    end
end

function regadd!(lhs::ArrayReg{B}, rhs::ArrayReg{B}) where {B}
    lhs.state .+= rhs.state
    lhs
end

function regsub!(lhs::ArrayReg{B}, rhs::ArrayReg{B}) where {B}
    lhs.state .-= rhs.state
    lhs
end

function regadd!(
    lhs::ArrayReg{B,T1,<:Transpose},
    rhs::ArrayReg{B,T2,<:Transpose},
) where {B,T1,T2}
    lhs.state.parent .+= rhs.state.parent
    lhs
end

function regsub!(
    lhs::ArrayReg{B,T1,<:Transpose},
    rhs::ArrayReg{B,T2,<:Transpose},
) where {B,T1,T2}
    lhs.state.parent .-= rhs.state.parent
    lhs
end

function regscale!(reg::ArrayReg{B,T1,<:Transpose}, x) where {B,T1}
    reg.state.parent .*= x
    reg
end

function regscale!(reg::ArrayReg{B}, x) where {B,T1}
    reg.state .*= x
    reg
end

# *, /
for op in [:*, :/]
    @eval function Base.$op(lhs::RT, rhs::Number) where {B,RT<:ArrayReg{B}}
        ArrayReg{B}($op(state(lhs), rhs))
    end

    @eval function Base.$op(lhs::RT, rhs::Number) where {B,RT<:AdjointArrayReg{B}}
        r = $op(parent(lhs), rhs')
        return adjoint(r)
    end

    if op == :*
        @eval function Base.$op(lhs::Number, rhs::RT) where {B,RT<:ArrayReg{B}}
            ArrayReg{B}(lhs * state(rhs))
        end

        @eval function Base.$op(lhs::Number, rhs::RT) where {B,RT<:AdjointArrayReg{B}}
            r = lhs' * parent(rhs)
            return adjoint(r)
        end
    end
end

for AT in [:ArrayReg, :AdjointArrayReg]
    @eval function Base.:(==)(lhs::$AT, rhs::$AT)
        state(lhs) == state(rhs)
    end

    @eval function Base.isapprox(lhs::$AT, rhs::$AT; kw...)
        isapprox(state(lhs), state(rhs); kw...)
    end
end

function Base.:*(bra::AdjointArrayReg{1}, ket::ArrayReg{1})
    if nremain(bra) == nremain(ket)
        return dot(relaxedvec(parent(bra)), relaxedvec(ket))
    elseif nremain(bra) == 0 # <s|active> |remain>
        return ArrayReg{1}(state(bra) * state(ket))
    else
        error(
            "partially contract ⟨bra|ket⟩ is not supported, expect ⟨bra| to be fully actived. nactive(bra)/nqubits(bra)=$(nactive(bra))/$(nqubits(bra))",
        )
    end
end

Base.:*(bra::AdjointArrayReg{B}, ket::ArrayReg{B}) where {B} = bra .* ket
function Base.:*(
    bra::AdjointArrayReg{B,T1,<:Transpose},
    ket::ArrayReg{B,T2,<:Transpose},
) where {B,T1,T2}
    if nremain(bra) == nremain(ket) == 0 # all active
        A, C = parent(state(parent(bra))), parent(state(ket))
        res = zeros(eltype(promote_type(T1, T2)), B)
        #return mapreduce((x, y) -> conj(x) * y, +, ; dims=2)
        for j = 1:size(A, 2)
            for i = 1:size(A, 1)
                @inbounds res[i] += conj(A[i, j]) * C[i, j]
            end
        end
        res
    elseif nremain(bra) == 0 # <s|active> |remain>
        bra .* ket
    else
        error(
            "partially contract ⟨bra|ket⟩ is not supported, expect ⟨bra| to be fully actived. nactive(bra)/nqubits(bra)=$(nactive(bra))/$(nqubits(bra))",
        )
    end
end

# broadcast
broadcastable(r::ArrayRegOrAdjointArrayReg{1}) = Ref(r)
broadcastable(r::ArrayRegOrAdjointArrayReg{B}) where {B} = (each for each in r)


text:
using LinearAlgebra:
    BlasReal,
    QRCompactWY,
    Diagonal,
    Hermitian,
    LowerTriangular,
    UpperTriangular,
    transpose,
    cholesky,
    cholesky!,
    qr,
    mul!,
    inv,
    diag,
    diagm,
    ldiv!

#################
modelmatrix(p::LinPred) = p.X
vcov(p::LinPred, wt::AbstractVector) =
    inv(Hermitian(float(Matrix(modelmatrix(p)' * (wt .* modelmatrix(p))))))
projectionmatrix(p::LinPred, wt::AbstractVector) =
    Hermitian(modelmatrix(p) * vcov(p, wt) * modelmatrix(p)' .* wt)

modelmatrix(m::RobustLinearModel) = modelmatrix(m.pred)

vcov(m::RobustLinearModel) = vcov(m.pred, workingweights(m.resp))

"""
    projectionmatrix(m::RobustLinearModel)

The robust projection matrix from the predictor: X (X' W X)⁻¹ X' W
"""
projectionmatrix(m::RobustLinearModel) = projectionmatrix(m.pred, workingweights(m.resp))

function leverage_weights(m::RobustLinearModel)
    w = weights(m.resp)
    v = inv(Hermitian(float(modelmatrix(m)' * (w .* modelmatrix(m)))))
    h = diag(Hermitian(modelmatrix(m) * v * modelmatrix(m)' .* w))
    sqrt.(1 .- h)
end

###
### From GLM, for information
###
#"""
#    linpred!(out, p::LinPred, f::Real=1.0)
#Overwrite `out` with the linear predictor from `p` with factor `f`
#The effective coefficient vector, `p.scratchbeta`, is evaluated as `p.beta0 .+ f * p.delbeta`,
#and `out` is updated to `p.X * p.scratchbeta`
#"""
#function linpred!(out, p::LinPred, f::Real=1.)
#    mul!(out, p.X, iszero(f) ? p.beta0 : broadcast!(muladd, p.scratchbeta, f, p.delbeta, p.beta0))
#end

#"""
#    linpred(p::LinPred, f::Read=1.0)
#Return the linear predictor `p.X * (p.beta0 .+ f * p.delbeta)`
#"""
#linpred(p::LinPred, f::Real=1.) = linpred!(Vector{eltype(p.X)}(undef, size(p.X, 1)), p, f)

#"""
#    installbeta!(p::LinPred, f::Real=1.0)
#Install `pbeta0 .+= f * p.delbeta` and zero out `p.delbeta`.  Return the updated `p.beta0`.
#"""
#function installbeta!(p::LinPred, f::Real=1.)
#    beta0 = p.beta0
#    delbeta = p.delbeta
#    @inbounds for i = eachindex(beta0,delbeta)
#        beta0[i] += delbeta[i]*f
#        delbeta[i] = 0
#    end
#    beta0
#end

"""
    SparsePredChol{T<:BlasReal} <: LinPred

A LinPred type with a sparse Cholesky factorization of X'X

# Members

- `X`: model matrix of size n×p with n ≥ p. Should be full column rank.
- `beta0`: base coefficient vector of length p
- `delbeta`: increment to coefficient vector, also of length p
- `scratchbeta`: scratch vector of length p, used in [`linpred!`](@ref) method
- `chol`: a Cholesky object created from X'X, possibly using row weights.
"""
SparsePredChol

"""
    DensePredCG

A `LinPred` type with Conjugate Gradient and a dense `X`

# Members

- `X`: Model matrix of size `n` × `p` with `n ≥ p`.  Should be full column rank.
- `beta0`: base coefficient vector of length `p`
- `delbeta`: increment to coefficient vector, also of length `p`
- `scratchbeta`: scratch vector of length `p`, used in [`linpred!`](@ref) method
"""
mutable struct DensePredCG{T<:BlasReal} <: DensePred
    X::Matrix{T}                  # model matrix
    beta0::Vector{T}              # base coefficient vector
    delbeta::Vector{T}            # coefficient increment
    scratchbeta::Vector{T}
    scratchm1::Matrix{T}
    scratchr1::Vector{T}
    function DensePredCG{T}(X::Matrix{T}, beta0::Vector{T}) where {T}
        n, p = size(X)
        length(beta0) == p || throw(DimensionMismatch("length(β0) ≠ size(X,2)"))
        new{T}(X, beta0, zeros(T, p), zeros(T, p), zeros(T, (n, p)), zeros(T, n))
    end
    function DensePredCG{T}(X::Matrix{T}) where {T}
        n, p = size(X)
        new{T}(X, zeros(T, p), zeros(T, p), zeros(T, p), zeros(T, (n, p)), zeros(T, n))
    end
end
DensePredCG(X::Matrix, beta0::Vector) = DensePredCG{eltype(X)}(X, beta0)
DensePredCG(X::Matrix{T}) where {T} = DensePredCG{T}(X, zeros(T, size(X, 2)))
convert(::Type{DensePredCG{T}}, X::Matrix{T}) where {T} =
    DensePredCG{T}(X, zeros(T, size(X, 2)))

"""
    delbeta!(p::LinPred, r::Vector)

Evaluate and return `p.delbeta` the increment to the coefficient vector from residual `r`
"""
function delbeta!(p::DensePredCG{T}, r::AbstractVector{T}) where {T<:BlasReal}
    lsqr!(p.delbeta, p.X, r; log=false)
    return p
end

cgpred(X::StridedMatrix) = DensePredCG(X)

function delbeta!(
    p::DensePredCG{T},
    r::AbstractVector{T},
    wt::AbstractVector{T},
) where {T<:BlasReal}
    ## Use views, do not create new objects
    # scr = transpose(broadcast!(*, p.scratchm1, wt, p.X))
    # cg!(p.delbeta,
    #     Hermitian(mul!(p.scratchm2, scr, p.X), :U),
    #     mul!(p.scratchbeta, scr, r)
    # )
    sqwt = copyto!(p.scratchr1, sqrt.(wt))
    scr = broadcast!(*, p.scratchm1, sqwt, p.X)
    lsqr!(p.delbeta, scr, broadcast!(*, p.scratchr1, sqwt, r); log=false)
    p
end

"""
    SparsePredCG

A `LinPred` type with Conjugate Gradient and a sparse `X`

# Members

- `X`: Model matrix of size `n` × `p` with `n ≥ p`.  Should be full column rank.
- `beta0`: base coefficient vector of length `p`
- `delbeta`: increment to coefficient vector, also of length `p`
- `scratchbeta`: scratch vector of length `p`, used in [`linpred!`](@ref) method
"""
mutable struct SparsePredCG{T,M<:SparseMatrixCSC} <: LinPred
    X::M                    # model matrix
    beta0::Vector{T}        # base vector for coefficients
    delbeta::Vector{T}      # coefficient increment
    scratchbeta::Vector{T}
    scratchm1::M
    scratchm2::M
    scratchr1::Vector{T}
end
function SparsePredCG(X::SparseMatrixCSC{T}) where {T}
    n, p = size(X)
    return SparsePredCG{eltype(X),typeof(X)}(
        X,
        zeros(T, p),
        zeros(T, p),
        zeros(T, p),
        similar(X),
        zeros(T, (p, p)),
        zeros(T, n),
    )
end

cgpred(X::SparseMatrixCSC) = SparsePredCG(X)

function delbeta!(
    p::SparsePredCG{T},
    r::AbstractVector{T},
    wt::AbstractVector{T}
) where {T}
    # scr = mul!(p.scratchm1, Diagonal(wt), p.X)
    # cg!(p.delbeta,
    #     Hermitian(mul!(transpose(scr), p.X), :U),
    #     mul!(p.scratchr1, transpose(scr), r)
    # )
    scr = transpose(broadcast!(*, p.scratchm1, wt, p.X))
    cg!(p.delbeta,
        Hermitian(mul!(p.scratchm2, scr, p.X), :U),
        mul!(p.scratchbeta, scr, r)
    )
    p
end

########
###     Ridge predictor
########

#function Base.float(::Type{S}, x::AbstractArray{T}) where {S<:AbstractFloat, T}
#    if !isconcretetype(T)
#        error("`float` not defined on abstractly-typed arrays;"*
#              " please convert to a more specific type")
#    end
#    convert(AbstractArray{typeof(convert(S, zero(T)))}, A)
#end

"""
    cat_ridge_matrix(X::AbstractMatrix{T}, λ::T, G::AbstractMatrix{T}) where {T}
        = vcat(X,  λ * G)

Construct the extended model matrix by vertically concatenating the regularizer to X.
"""
cat_ridge_matrix(X::AbstractMatrix{T}, λ::T, G::AbstractMatrix{T}) where {T} =
    vcat(X, λ * G)

"""
    RidgePred

Regularized predictor using ridge regression on the `p` features.

# Members

- `X`: model matrix
- `λ`: shrinkage parameter of the regularizer
- `G`: regularizer matrix of size p×p.
- `βprior`: regularizer prior of the coefficient values. Default to `zeros(p)`.
- `pred`: the non-regularized predictor using an extended model matrix.
- `pivot`: for `DensePredChol`, if the decomposition was pivoted.
- `scratchbeta`: scratch vector of length `p`, used in [`linpred!`](@ref) method
"""
mutable struct RidgePred{T<:BlasReal,M<:AbstractMatrix,P<:LinPred} <:
               AbstractRegularizedPred{T}
    X::M                    # model matrix
    sqrthalfλ::T            # sqrt of half of the shrinkage parameter λ
    G::M                    # regularizer matrix
    βprior::Vector{T}       # coefficients prior
    pred::P                 # predictor
    pivot::Bool             # pivoted decomposition
    scratchbeta::Vector{T}
    scratchy::Vector{T}
    scratchwt::Vector{T}
end

function RidgePred(
    ::Type{P},
    X::M,
    λ::T,
    G::M,
    βprior::Vector{T},
    pivot::Bool=false,
) where {M<:Union{SparseMatrixCSC{T},Matrix{T}},P<:LinPred} where {T<:BlasReal}
    λ >= 0 || throw(DomainError(λ, "the shrinkage parameter should be non-negative"))

    m1, m2 = size(G)
    m1 == m2 || throw(DimensionMismatch("the regularization matrix should be square"))
    n, m = size(X)
    m1 == m || throw(
        DimensionMismatch(
            "the regularization matrix should of size p×p," *
            " with p the number of predictor in matrix X:" *
            " size(G)=$(size(G)) != size(X, 2)=$(m)",
        ),
    )
    ll = size(βprior, 1)
    ll in (0, m) || throw(DimensionMismatch("length of βprior is $ll, must be $m or 0"))

    sqrthalfλ = √(λ / 2)
    pred = if P == DensePredChol
        P(cat_ridge_matrix(X, sqrthalfλ, G), pivot)
    else
        P(cat_ridge_matrix(X, sqrthalfλ, G))
    end

    return RidgePred{T,typeof(X),typeof(pred)}(
        X,
        sqrthalfλ,
        G,
        (ll == 0) ? zeros(T, m) : βprior,
        pred,
        pivot,
        zeros(T, m),
        zeros(T, n + m),
        ones(T, n + m),
    )
end

function postupdate_λ!(r::RidgePred)
    n, m = size(r.X)
    # Update the extended model matrix with the new value
    GG = r.sqrthalfλ * r.G
    @views r.pred.X[n+1:n+m, :] .= GG
    if isa(r.pred, DensePredChol)
        # Recompute the cholesky decomposition
        X = r.pred.X
        F = Hermitian(float(X'X))
        T = eltype(F)
        r.pred.chol =
            r.pivot ? cholesky!(F, Val(true), tol=-one(T), check=false) : cholesky!(F)
    elseif isa(r.pred, SparsePredChol)
        # Update Xt
        @views r.pred.Xt[:, n+1:n+m] .= GG'
    end
end

function Base.getproperty(r::RidgePred, s::Symbol)
    if s ∈ (:λ, :lambda)
        2 * (r.sqrthalfλ)^2
    elseif s ∈ (:beta0, :delbeta)
        getproperty(r.pred, s)
    else
        getfield(r, s)
    end
end

function Base.setproperty!(r::RidgePred, s::Symbol, v)
    if s ∈ (:λ, :lambda)
        v >= 0 || throw(DomainError(λ, "the shrinkage parameter should be non-negative"))
        # Update the square root value
        r.sqrthalfλ = √(v / 2)
        postupdate_λ!(r)
    elseif s ∈ (:sqrthalfλ,)
        setfield!(r, s, v)
    else
        error(
            "cannot set any property of RidgePred except" *
            " the shrinkage parameter λ (lambda): $(s)",
        )
        # setfield!(r, s, v)
    end
end

function Base.propertynames(r::RidgePred, private=false)
    if private
        (
            :X,
            :λ,
            :lambda,
            :G,
            :βprior,
            :pred,
            :beta0,
            :delbeta,
            :pivot,
            :sqrthalfλ,
            :scratchbeta,
            :scratchy,
            :scratchwt,
        )
    else
        (:X, :λ, :lambda, :G, :βprior, :pred, :beta0, :delbeta)
    end
end

function cgpred(
    X::StridedMatrix,
    λ::Real,
    G::AbstractMatrix,
    βprior::AbstractVector=zeros(eltype(X), size(X, 2)),
)
    T = eltype(X)
    RidgePred(DensePredCG, X, Base.convert(T, λ), Matrix{T}(G), Vector{T}(βprior))
end

function cholpred(
    X::StridedMatrix,
    λ::Real,
    G::AbstractMatrix,
    βprior::AbstractVector=zeros(eltype(X), size(X, 2)),
    pivot::Bool=false,
)
    T = eltype(X)
    RidgePred(DensePredChol, X, Base.convert(T, λ), Matrix{T}(G), Vector{T}(βprior), pivot)
end

function cgpred(
    X::SparseMatrixCSC,
    λ::Real,
    G::AbstractMatrix,
    βprior::AbstractVector=zeros(eltype(X), size(X, 2)),
)
    T = eltype(X)
    RidgePred(SparsePredCG, X, Base.convert(T, λ), SparseMatrixCSC{T}(G), Vector{T}(βprior))
end

function cholpred(
    X::SparseMatrixCSC,
    λ::Real,
    G::AbstractMatrix,
    βprior::AbstractVector=zeros(eltype(X), size(X, 2)),
)
    T = eltype(X)
    RidgePred(
        SparsePredChol,
        X,
        Base.convert(T, λ),
        SparseMatrixCSC{T}(G),
        Vector{T}(βprior),
    )
end

cgpred(X::SparseMatrixCSC, λ, G::AbstractVector, βprior::AbstractVector) =
    cgpred(X, λ, spdiagm(0 => G), βprior)

cholpred(X::SparseMatrixCSC, λ, G::AbstractVector, βprior::AbstractVector) =
    cholpred(X, λ, spdiagm(0 => G), βprior)

cgpred(X::StridedMatrix, λ, G::AbstractVector, βprior::AbstractVector) =
    cgpred(X, λ, diagm(0 => G), βprior)

cholpred(X::StridedMatrix, λ, G::AbstractVector, βprior::AbstractVector, args...) =
    cholpred(X, λ, diagm(0 => G), βprior, args...)

function resetβ0!(p::RidgePred{T}) where {T<:BlasReal}
    beta0 = p.beta0
    delbeta = p.delbeta
    @inbounds for i in eachindex(beta0, delbeta)
        beta0[i] = 0
        delbeta[i] = 0
    end
    p
end

function delbeta!(
    p::RidgePred{T},
    r::AbstractVector{T},
    wt::AbstractVector{T},
) where {T<:BlasReal}
    n, m = size(p.X)
    # Fill response
    copyto!(p.scratchy, r)
    # yprior = sqrt(p.λ/2) * p.G * (p.βprior - p.pred.beta0)
    broadcast!(-, p.scratchbeta, p.βprior, p.pred.beta0)
    @views mul!(p.scratchy[n+1:end], p.G, p.scratchbeta, p.sqrthalfλ, 0)

    # Fill weights
    copyto!(p.scratchwt, wt)

    # Compute Δβₙ from (XᵀWₙX + λGᵀG)⁻¹ (XᵀWₙrₙ + λGᵀG(βprior - βₙ))
    delbeta!(p.pred, p.scratchy, p.scratchwt)

    p
end

extendedmodelmatrix(p::RidgePred) = p.pred.X

function extendedweights(p::RidgePred{T}, wt::AbstractVector) where {T<:BlasReal}
    n, m = size(p.X)
    k = length(wt)
    if k == 0
        return ones(T, n + m)
    elseif k == n
        return vcat(Vector{T}(wt), ones(T, m))
    elseif k == n + m
        return wt
    else
        throw(DimensionMismatch("weights for RidgePred should be of size 0, n or n+m"))
    end
end

modelmatrix(p::RidgePred) = p.X

function vcov(p::RidgePred, wt::AbstractVector)
    pp = extendedmodelmatrix(p)' * (extendedweights(p, wt) .* extendedmodelmatrix(p))
    return inv(Hermitian(float(Matrix(pp))))
end

function projectionmatrix(p::RidgePred, wt::AbstractVector)
    wwt = extendedweights(p, wt)
    Hermitian(extendedmodelmatrix(p) * vcov(p, wwt) * extendedmodelmatrix(p)' .* wwt)
end

dof(m::RobustLinearModel{T,R,L}) where {T,R,L<:AbstractRegularizedPred} =
    tr(projectionmatrix(m.pred, workingweights(m.resp)))


text:
compartment = SBML.MathIdent("compartment")
k1 = SBML.MathIdent("k1")
k2 = SBML.MathIdent("k2")
reaction1_k = SBML.MathIdent("reaction1_k")
S1 = SBML.MathIdent("S1")

sbmlfiles = [
    # sbml_test_suite case 00001
    (
        joinpath(@__DIR__, "data", "00001-sbml-l3v2.xml"),
        "https://raw.githubusercontent.com/sbmlteam/sbml-test-suite/master/cases/semantic/00001/00001-sbml-l3v2.xml",
        "3e9a8bfe27343459d7a1e462fc3fd4eb2b01dc7b32af1db06a98f366287da01a",
        x -> nothing,
        1,
        ("reaction1", SBML.MathApply("*", [SBML.MathApply("*", [compartment, k1]), S1])),
        ("S1", 0.00015),
    ),
    # case 00001 in older level and version
    (
        joinpath(@__DIR__, "data", "00001-sbml-l2v1.xml"),
        "https://raw.githubusercontent.com/sbmlteam/sbml-test-suite/master/cases/semantic/00001/00001-sbml-l2v1.xml",
        "71a145c58b08e475d76bdec644589b2a55b5c5c2fee218274c91677c0f30b508",
        SBML.set_level_and_version(3, 2),
        1,
        ("reaction1", SBML.MathApply("*", [SBML.MathApply("*", [compartment, k1]), S1])),
        ("S1", 0.00015),
    ),
    # case 00057 with localParameters
    (
        joinpath(@__DIR__, "data", "00057-sbml-l3v2.xml"),
        "https://raw.githubusercontent.com/sbmlteam/sbml-test-suite/master/cases/semantic/00057/00057-sbml-l3v2.xml",
        "3e84e19cebbb79eea879847f541b1d22db6eb239f1f070ef4609f04c77688659",
        SBML.libsbml_convert("promoteLocalParameters"),
        2,
        (
            "reaction1",
            SBML.MathApply("*", [SBML.MathApply("*", [compartment, reaction1_k]), S1]),
        ),
        ("S1", 0.0003),
    ),
    # case 00025 with functionDefinition
    (
        joinpath(@__DIR__, "data", "00025-sbml-l3v2.xml"),
        "https://raw.githubusercontent.com/sbmlteam/sbml-test-suite/master/cases/semantic/00025/00025-sbml-l3v2.xml",
        "d3231ae3858d9e5dca1b106aa7b106a0caee9e6967ef8413de6b9acde9171c3e",
        SBML.libsbml_convert("expandFunctionDefinitions"),
        1,
        ("reaction1", SBML.MathApply("*", [compartment, SBML.MathApply("*", [k1, S1])])),
        ("S1", 0.0015),
    ),
    # case 00036 with initialAssignment
    (
        joinpath(@__DIR__, "data", "00037-sbml-l3v2.xml"),
        "https://raw.githubusercontent.com/sbmlteam/sbml-test-suite/master/cases/semantic/00037/00037-sbml-l3v2.xml",
        "074f0caeaf2cdc390967bfdf06d80ccad519c648df7f421dc4e9c69e71551dff",
        SBML.libsbml_convert("expandInitialAssignments"),
        2,
        ("reaction1", SBML.MathApply("*", [SBML.MathApply("*", [compartment, k2]), S1])),
        ("S1", 0.01125),
    ),
]

@testset "Loading of models from sbml_test_suite" begin
    for (sbmlfile, url, hash, converter, expected_par, expected_rxn, expected_u0) in
        sbmlfiles
        if !isfile(sbmlfile)
            Downloads.download(url, sbmlfile)
        end

        cksum = bytes2hex(sha256(open(sbmlfile)))
        if cksum != hash
            @warn "The downloaded model `$sbmlfile' seems to be different from the expected one. Tests will likely fail." cksum
        end

        @testset "Loading of $sbmlfile" begin
            mdl = readSBML(sbmlfile, converter)

            @test typeof(mdl) == Model

            @test length(mdl.parameters) == expected_par

            id, math = expected_rxn
            @test isequal(repr(mdl.reactions[id].kinetic_math), repr(math))

            id, ic = expected_u0
            if basename(sbmlfile) == "00037-sbml-l3v2.xml"
                # When expanding initial assignments with libsbml, the initial amount
                # becomes empty.
                @test_broken mdl.species[id].initial_amount[1] == ic
            else
                @test mdl.species[id].initial_amount[1] == ic
            end
        end
    end
end


text:
export HCircle, get_center, get_radius, points_on_circle
export RandomHCircle, circumference, area

"""
`HCircle(P::HPoint,r::Real)` creates a new hyperbolic circle
centered at `P` with radius `r`.
"""
struct HCircle <: HObject
    ctr::HPoint
    rad::Float64
    attr::Dict{Symbol,Any}
    function HCircle(P::HPoint,r::Real)
        @assert r>0 "radius must be positive"
        C = new(P,r,Dict{Symbol,Any}())
        set_color(C)
        set_thickness(C)
        set_line_style(C)
        return C
    end
end
HCircle(C::HCircle) = HCircle(C.ctr, C.rad)  # copy constructor

"""
`HCircle(A,B,C)` creates a circle that includes the three given points.
"""
function HCircle(A::HPoint, B::HPoint, C::HPoint)
    @assert !collinear(A,B,C) "The three points must be nonlinear"
    L1 = bisector(A,B)
    L2 = bisector(A,C)
    @assert meet_check(L1,L2) "There is no circle containing these three points"
    P = meet(L1,L2)
    r = dist(A,P)
    return HCircle(P,r)
end

"""
`RandomHCircle()` creates a random circle.
"""
RandomHCircle() = HCircle(RandomHPoint(), -log(rand()) )

"""
`get_center(C::HCircle)` returns the center of the circle.
"""
get_center(C::HCircle) = C.ctr

"""
`get_radius(C::HCircle)` returns the radius of the circle.
"""
get_radius(C::HCircle) = C.rad

function show(io::IO, C::HCircle)
    print(io,"HCircle($(C.ctr),$(C.rad))")
end

"""
`points_on_circle(C::HCircle)` returns a 3-tuple of points
that lie on the circle `C`.
"""
function points_on_circle(C::HCircle)
    P = get_center(C)
    f = inv(move2zero(P))
    r = get_radius(C)
    a = HPoint(r,0)
    b = HPoint(r,2pi/3)
    c = HPoint(r,4pi/3)
    return (f(a), f(b), f(c))
end

HCircle(r::Real = 1) = HCircle(HPoint(), r)

function(==)(C::HCircle, CC::HCircle)
    if abs(C.rad - CC.rad) > THRESHOLD*eps(1.0)
        return false
    end
    return C.ctr == CC.ctr
end

"""
`circumference(C::HCircle)` returns the circumference of the circle.
"""
function circumference(C::HCircle)
    r = get_radius(C)
    return 2*pi*sinh(r)
end

"""
`area(C::HCircle)` returns the area of the circle.
"""
function area(C::HCircle)
    r = get_radius(C)
    return 4*pi*(sinh(r/2))^2
end

function in(P::HPoint, C::HCircle)::Bool
    d = dist(P, C.ctr)
    return d <= C.rad + THRESHOLD*eps(1.0)   # add a bit of slop
end


text:
# Serialization is imported only if cache is used

function write_cache(doc::WeaveDoc, cache_path)
    cache_dir = joinpath(doc.cwd, cache_path)
    isdir(cache_dir) || mkpath(cache_dir)
    open(joinpath(cache_dir, doc.basename * ".cache"), "w") do io
        Serialization.serialize(io, doc)
    end
    return nothing
end

function read_cache(doc::WeaveDoc, cache_path)
    name = joinpath(doc.cwd, cache_path, doc.basename * ".cache")
    isfile(name) || return nothing
    open(name, "r") do io
        doc = Serialization.deserialize(io)
    end
    return doc
end

function restore_chunk(chunk::CodeChunk, cached)
    chunks = filter(
        x -> x.number == chunk.number && string(typeof(x)) == "Weave.CodeChunk",
        cached.chunks,
    )

    # Chunk types, don't match after loading. Fix by constructing chunks
    # from loaded content
    new_chunks = []
    for c in chunks
        newc = CodeChunk(c.content, c.number, c.start_line, c.optionstring, c.options)
        newc.figures = c.figures
        newc.result = c.result
        newc.output = c.output
        newc.rich_output = c.rich_output
        push!(new_chunks, newc)
    end
    return new_chunks
end

# Restore inline code
function restore_chunk(chunk::DocChunk, cached::WeaveDoc)
    # Get chunk from cached doc
    c_chunk = filter(x -> x.number == chunk.number && isa(x, DocChunk), cached.chunks)
    isempty(c_chunk) && return chunk
    c_chunk = c_chunk[1]

    # Collect cached code
    c_inline = filter(x -> isa(x, InlineCode), c_chunk.content)
    isempty(c_inline) && return chunk

    # Restore cached results for Inline code
    n = length(chunk.content)
    for i = 1:n
        if isa(chunk.content[i], InlineCode)
            ci = filter(x -> x.number == chunk.content[i].number, c_inline)
            isempty(ci) && continue
            chunk.content[i].output = ci[1].output
            chunk.content[i].rich_output = ci[1].rich_output
            chunk.content[i].figures = ci[1].figures
        end
    end
    return chunk
end


text:
using Rimu
using Test
using Rimu.DictVectors: Initiator, SimpleInitiator, CoherentInitiator
using Rimu.StochasticStyles: IsStochastic2Pop, Bernoulli, WithoutReplacement, IsExplosive
using Rimu.StochasticStyles: DoubleOrNothing, ThresholdCompression
using Rimu.StatsTools
using Rimu.ConsistentRNG: seedCRNG!
using Rimu.RMPI
using KrylovKit
using Suppressor
using Statistics

@testset "lomc!/QMCState" begin
    @testset "Setting laststep" begin
        add = BoseFS{5,2}((2,3))
        H = HubbardReal1D(add; u=0.1)
        dv = DVec(add => 1; style=IsStochasticInteger())

        df, state = lomc!(H, copy(dv); laststep=100)
        @test size(df, 1) == 100
        @test state.replicas[1].params.step == 100

        df, state = lomc!(H, copy(dv); laststep=200)
        @test size(df, 1) == 200
        @test state.replicas[1].params.step == 200

        df, state = lomc!(H, copy(dv); laststep=13)
        @test size(df, 1) == 13
        @test state.replicas[1].params.step == 13

        state.laststep = 100
        df = lomc!(state, df).df
        @test size(df, 1) == 100

        state.step = 0
        df = lomc!(state, df).df
        @test size(df, 1) == 200
        @test df.steps == [1:100; 1:100]
    end

    @testset "Setting walkernumber" begin
        add = BoseFS{2,5}((0,0,2,0,0))
        H = HubbardMom1D(add; u=0.5)
        dv = DVec(add => 1; style=IsStochasticWithThreshold(1.0))

        s_strat = DoubleLogUpdate(ζ=0.05, ξ=0.05^2/4, targetwalkers=100)
        v = copy(dv)
        walkers = lomc!(H, v; s_strat, laststep=1000).df.norm
        @test median(walkers) ≈ 100 rtol=0.1
        s_strat = LogUpdate(0.05)
        walkers = lomc!(H, v; s_strat, laststep=1000).df.norm # continuation run
        @test median(walkers) > 10 # essentially just test that it does not error

        s_strat = DoubleLogUpdate(ζ=0.05, ξ=0.05^2/4, targetwalkers=200)
        walkers = lomc!(H, copy(dv); s_strat, laststep=1000).df.norm
        @test median(walkers) ≈ 200 rtol=0.1

        s_strat = DoubleLogUpdate(ζ=0.05, ξ=0.05^2/4, targetwalkers=1000)
        walkers = lomc!(H, copy(dv); s_strat, laststep=1000).df.norm
        @test median(walkers) ≈ 1000 rtol=0.1
    end

    @testset "Replicas" begin
        add = near_uniform(BoseFS{5,15})
        H = HubbardReal1D(add)
        dv = DVec(add => 1, style=IsDynamicSemistochastic())

        @testset "NoStats" begin
            df, state = lomc!(H, dv; replica=NoStats(1))
            @test state.replica == NoStats(1)
            @test length(state.replicas) == 1
            @test "shift" ∈ names(df)
            @test "shift_1" ∉ names(df)

            df, state = lomc!(H, dv; replica=NoStats(3))
            @test state.replica == NoStats(3)
            @test length(state.replicas) == 3
            @test df.shift_1 ≠ df.shift_2 && df.shift_2 ≠ df.shift_3
            @test "shift_4" ∉ names(df)
        end

        @testset "AllOverlaps" begin
            # column names are of the form c{i}_dot_c{j} and c{i}_Op{k}_c{j}.
            num_stats(df) = length(filter(x -> match(r"^c[0-9]", x) ≠ nothing, names(df)))

            # No operator: N choose 2 reports.
            df, _ = lomc!(H, dv; replica=AllOverlaps(4))
            @test num_stats(df) == binomial(4, 2)
            df, _ = lomc!(H, dv; replica=AllOverlaps(5))
            @test num_stats(df) == binomial(5, 2)

            # One operator: 2 * N choose 2 reports.
            df, _ = lomc!(H, dv; replica=AllOverlaps(4, H))
            @test num_stats(df) == 2 * binomial(4, 2)
            df, _ = lomc!(H, dv; replica=AllOverlaps(5, H))
            @test num_stats(df) == 2 * binomial(5, 2)

            # Two operators: 3 * N choose 2 reports.
            df, _ = lomc!(H, dv; replica=AllOverlaps(2, (GutzwillerSampling(H, 1), H)))
            @test num_stats(df) == 3 * binomial(2, 2)
            df, _ = lomc!(H, dv; replica=AllOverlaps(7, (GutzwillerSampling(H, 1), H)))
            @test num_stats(df) == 3 * binomial(7, 2)

            # Complex operator
            v = DVec(1 => 1)
            G = MatrixHamiltonian(rand(5, 5))
            O = MatrixHamiltonian(rand(ComplexF64, 5, 5))
            df, _ = lomc!(G, v, replica=AllOverlaps(2, O))
            @test df.c1_dot_c2 isa Vector{ComplexF64}
            @test df.c1_Op1_c2 isa Vector{ComplexF64}

            # MPIData
            df, _ = lomc!(H, MPIData(dv); replica=AllOverlaps(4, H))
            @test num_stats(df) == 2 * binomial(4, 2)
            df, _ = lomc!(H, MPIData(dv); replica=AllOverlaps(5, H))
            @test num_stats(df) == 2 * binomial(5, 2)
        end
    end

    @testset "Dead population" begin
        add = BoseFS{5,2}((2,3))
        H = HubbardReal1D(add; u=20)
        dv = DVec(add => 1; style=IsStochasticInteger())

        # Only population is dead.
        params = RunTillLastStep(shift = 0.0)
        df = @suppress_err lomc!(H, copy(dv); params, laststep=100).df
        @test size(df, 1) < 100

        # population does not die with sensible default shift
        df = lomc!(H, copy(dv); laststep=100).df
        @test size(df, 1) == 100

        # Populations in replicas are dead.
        params = RunTillLastStep(shift = 0.0)
        df = @suppress_err lomc!(H, copy(dv); params, laststep=100, replica=NoStats(5)).df
        @test size(df, 1) < 100
    end

    @testset "Default DVec" begin
        add = BoseFS{5,2}((2,3))
        H = HubbardReal1D(add; u=20)
        df, state = lomc!(H; laststep=100)
        @test StochasticStyle(state.replicas[1].v) isa IsStochasticInteger

        df, state = lomc!(H; laststep=100, style = IsDeterministic())
        @test StochasticStyle(state.replicas[1].v) isa IsDeterministic
    end

    @testset "ShiftStrategy" begin
        add = BoseFS{5,2}((2,3))
        H = HubbardReal1D(add; u=20)

        # DontUpdate
        s_strat = DontUpdate(targetwalkers = 100)
        df = lomc!(H; s_strat, laststep=100).df
        @test size(df, 1) < 100 # finish early without error

        # LogUpdateAfterTargetWalkers
        s_strat = LogUpdateAfterTargetWalkers(targetwalkers = 100)
        df, state  = lomc!(H; s_strat, laststep=100)
        @test size(df, 1) == 100
        @test df.shiftMode[end] # finish in variable shift mode
        @test df.norm[end] > 100

        # LogUpdate
        v = state.replicas[1].v
        params = state.replicas[1].params
        s_strat = LogUpdate()
        df = lomc!(H, v; df, params, s_strat, laststep=200).df
        @test size(df, 1) == 200
        @test 500 > df.norm[end] > 100

        # DoubleLogUpdateAfterTargetWalkers
        s_strat = DoubleLogUpdateAfterTargetWalkers(targetwalkers = 100)
        df, state  = lomc!(H; s_strat, laststep=100)
        @test size(df, 1) == 100
        @test df.shiftMode[end] # finish in variable shift mode
        @test df.norm[end] > 100

        # test unexported strategies
        # DoubleLogSumUpdate
        s_strat = Rimu.DoubleLogSumUpdate(targetwalkers = 100)
        df, state  = lomc!(H; s_strat, laststep=100)
        @test size(df, 1) == 100
        @test df.shiftMode[end] # finish in variable shift mode

        # TripleLogUpdate
        s_strat = Rimu.TripleLogUpdate(targetwalkers = 100)
        df, state  = lomc!(H; s_strat, laststep=100)
        @test size(df, 1) == 100
        @test df.shiftMode[end] # finish in variable shift mode

        # DoubleLogProjected
        s_strat = Rimu.DoubleLogProjected(target = 100.0, projector=UniformProjector())
        df, state  = lomc!(H; s_strat, laststep=100)
        @test size(df, 1) == 100
        @test df.shiftMode[end] # finish in variable shift mode

    end

    @testset "deprecated" begin
        @test_throws ErrorException DelayedLogUpdateAfterTargetWalkers()
        @test_throws ErrorException HistoryLogUpdate()
        @test_throws ErrorException DelayedLogUpdate()
        @test_throws ErrorException DelayedDoubleLogUpdateAfterTW()
        @test_throws ErrorException DoubleLogUpdateAfterTargetWalkersSwitch()
        @test_throws ErrorException DelayedDoubleLogUpdate()
    end

    @testset "Setting `maxlength`" begin
        add = BoseFS{15,10}((0,0,0,0,0,15,0,0,0,0))
        H = HubbardMom1D(add; u=6.0)
        dv = DVec(add => 1; style=IsDynamicSemistochastic())

        seedCRNG!(1336)

        df = @suppress_err lomc!(H, copy(dv); maxlength=10, dτ=1e-4).df
        @test all(df.len[1:end-1] .≤ 10)
        @test df.len[end] > 10

        df, state = @suppress_err lomc!(H, copy(dv); maxlength=10, dτ=1e-4, replica=NoStats(6))
        @test all(df.len_1[1:end-1] .≤ 10)
        @test all(df.len_2[1:end-1] .≤ 10)
        @test all(df.len_3[1:end-1] .≤ 10)
        @test all(df.len_4[1:end-1] .≤ 10)
        @test all(df.len_5[1:end-1] .≤ 10)
        @test all(df.len_6[1:end-1] .≤ 10)

        state.maxlength += 1000
        df_cont = lomc!(state).df
        @test size(df_cont, 1) == 100 - size(df, 1)
    end

    @testset "Continuations" begin
        add = BoseFS{5,5}((1,1,1,1,1))
        H = HubbardReal1D(add; u=0.5)
        # Using Deterministic to get exact same result
        dv = DVec(add => 1.0, style=IsDeterministic())

        # Run lomc!, then change laststep and continue.
        df, state = lomc!(H, copy(dv); threading=false)
        state.laststep = 200
        df1 = lomc!(state, df).df

        # Run lomc! with laststep already set.
        df2 = lomc!(H, copy(dv); laststep=200).df

        @test df1.len ≈ df2.len
        @test df1.norm ≈ df2.norm
        @test df1.shift ≈ df2.shift
    end

    @testset "Reporting" begin
        add = BoseFS((1,2,1,1))
        H = HubbardReal1D(add; u=2)
        dv = DVec(add => 1, style=IsDeterministic())

        @testset "ReportDFAndInfo" begin
            r_strat = ReportDFAndInfo(k=5, i=20, io=devnull, writeinfo=true)
            df = lomc!(H, copy(dv); r_strat, laststep=100).df
            @test size(df, 1) == 20

            out = @capture_out begin
                r_strat = ReportDFAndInfo(k=5, i=20, io=stdout, writeinfo=true)
                lomc!(H, copy(dv); r_strat, laststep=100)
            end
            @test length(split(out, '\n')) == 6 # (last line is empty)
        end
        @testset "ReportToFile" begin
            # Clean up.
            rm("test-report.arrow"; force=true)
            rm("test-report-1.arrow"; force=true)
            rm("test-report-2.arrow"; force=true)

            r_strat = ReportToFile(filename="test-report.arrow", io=devnull, save_if=false)
            df = lomc!(H, copy(dv); r_strat, laststep=100).df
            @test !isfile("test-report.arrow")

            r_strat = ReportToFile(filename="test-report.arrow", io=devnull)
            df = lomc!(H, copy(dv); r_strat, laststep=100).df
            @test isempty(df)
            df1 = RimuIO.load_df("test-report.arrow")

            r_strat = ReportToFile(filename="test-report.arrow", io=devnull, chunk_size=5)
            df = lomc!(H, copy(dv); r_strat, laststep=100).df
            @test isempty(df)
            df2 = RimuIO.load_df("test-report-1.arrow")

            r_strat = ReportToFile(filename="test-report.arrow", io=devnull, return_df=true)
            df3 = lomc!(H, copy(dv); r_strat, laststep=100).df
            @test isempty(df)
            df4 = RimuIO.load_df("test-report-2.arrow")

            @test df1.shift ≈ df2.shift
            @test df2.norm ≈ df3.norm
            @test df3 == df4

            # Clean up.
            rm("test-report.arrow"; force=true)
            rm("test-report-1.arrow"; force=true)
            rm("test-report-2.arrow"; force=true)
        end
    end

    @testset "Post step" begin
        add = BoseFS((0, 0, 0, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0))
        H = HubbardMom1D(add; u=4)
        dv = DVec(add => 1)

        @testset "Projector, ProjectedEnergy" begin
            ConsistentRNG.seedCRNG!(1337)

            post_step = (
                Projector(p1=NormProjector()),
                Projector(p2=copy(dv)),
                ProjectedEnergy(H, dv),
                ProjectedEnergy(H, dv, vproj=:vproj2, hproj=:hproj2),
            )
            df, _ = lomc!(H, copy(dv); post_step)
            @test df.vproj == df.vproj2 == df.p2
            @test df.norm ≈ df.p1

            @test_throws ErrorException lomc!(
                H, dv; post_step=(Projector(a=dv), Projector(a=dv))
            )
            @test_throws ErrorException Projector(a=dv, b=dv)
            @test_throws ErrorException Projector()
        end

        @testset "SignCoherence" begin
            ConsistentRNG.seedCRNG!(1337)

            ref = eigsolve(H, dv, 1, :SR; issymmetric=true)[2][1]
            post_step = (SignCoherence(ref), SignCoherence(dv * -1, name=:single_coherence))
            df, _ = lomc!(H, copy(dv); post_step)
            @test df.coherence[1] == 1.0
            @test all(-1.0 .≤ df.coherence .≤ 1.0)
            @test all(in.(df.single_coherence, Ref((-1, 0, 1))))

            cdv = DVec(add => 1 + im)
            df, _ = lomc!(H, cdv; post_step)
            @test df.coherence isa Vector{ComplexF64}
        end

        @testset "WalkerLoneliness" begin
            ConsistentRNG.seedCRNG!(1337)

            post_step = WalkerLoneliness()
            df, _ = lomc!(H, copy(dv); post_step)
            @test df.loneliness[1] == 1
            @test all(1 .≥ df.loneliness .≥ 0)

            cdv = DVec(add => 1 + im)
            df, _ = lomc!(H, cdv; post_step)
            @test df.loneliness isa Vector{ComplexF64}
        end

        @testset "Timer" begin
            post_step = Rimu.Timer()
            time_before = time()
            df, _ = lomc!(H, copy(dv); post_step)
            time_after = time()
            @test df.time[1] > time_before
            @test df.time[end] < time_after
            @test issorted(df.time)
        end
    end

    @testset "Threading" begin
        add = BoseFS((0,0,0,0,10,0,0,0,0,0))
        H = HubbardMom1D(add)
        s_strat = DoubleLogUpdate(targetwalkers=10_000)
        laststep = 1000
        for dv in (DVec(add => 1), InitiatorDVec(add => 1), MPIData(DVec(add => 1)))
            @test_throws ErrorException lomc!(H; threading=:something, laststep, s_strat)

            df1, s1 = lomc!(
                H, deepcopy(dv); threading=:auto, laststep, s_strat
            )
            df2, s2 = lomc!(
                H, deepcopy(dv); threading=true, laststep, s_strat
            )
            df3, s3 = lomc!(
                H, deepcopy(dv); threading=false, laststep, s_strat
            )
            df4, s4 = lomc!(
                H, deepcopy(dv); threading=Rimu.NoThreading(), laststep, s_strat
            )
            df5, s5 = lomc!(
                H, deepcopy(dv); threading=Rimu.ThreadsThreading(), laststep, s_strat
            )
            df6, s6 = lomc!(
                H, deepcopy(dv); threading=Rimu.SplittablesThreading(), laststep, s_strat
            )
            df7, s7 = lomc!(
                H, deepcopy(dv); threading=Rimu.ThreadsXSumThreading(), laststep, s_strat
            )
            df8, s8 = lomc!(
                H, deepcopy(dv); threading=Rimu.ThreadsXForeachThreading(), laststep, s_strat
            )

            @test s1.threading == Rimu.SplittablesThreading()
            @test s2.threading == Rimu.SplittablesThreading()
            @test s3.threading == Rimu.NoThreading()

            # Check that working memory was selected correctly.
            N = Threads.nthreads()
            D = typeof(localpart(dv))
            @test s1.replicas[1].w isa NTuple{N,D}
            @test s2.replicas[1].w isa NTuple{N,D}
            @test s3.replicas[1].w isa D
            @test s4.replicas[1].w isa D
            @test s5.replicas[1].w isa NTuple{N,D}
            @test s6.replicas[1].w isa NTuple{N,D}
            @test s7.replicas[1].w isa NTuple{N,D}
            @test s8.replicas[1].w isa NTuple{N,D}

            # Check energies.
            E1, _ = mean_and_se(df1.shift[900:end])
            E2, _ = mean_and_se(df2.shift[900:end])
            E3, _ = mean_and_se(df3.shift[900:end])
            E4, _ = mean_and_se(df4.shift[900:end])
            E5, _ = mean_and_se(df5.shift[900:end])
            E6, _ = mean_and_se(df6.shift[900:end])
            E7, _ = mean_and_se(df7.shift[900:end])
            E8, _ = mean_and_se(df8.shift[900:end])

            @test E1 ≈ E2 rtol=0.1
            @test E2 ≈ E3 rtol=0.1
            @test E3 ≈ E4 rtol=0.1
            @test E4 ≈ E5 rtol=0.1
            @test E5 ≈ E6 rtol=0.1
            @test E6 ≈ E7 rtol=0.1
            @test E7 ≈ E8 rtol=0.1
        end
    end
end

@testset "Ground state energy estimates" begin
    for H in (
        HubbardReal1D(BoseFS((1,1,2))),
        BoseHubbardReal1D2C(BoseFS2C((1,2,2), (0,1,0))),
        BoseHubbardMom1D2C(BoseFS2C((0,1), (1,0))),
    )
        @testset "$H" begin
            dv = DVec(starting_address(H) => 2; style=IsDynamicSemistochastic())
            post_step = ProjectedEnergy(H, dv)

            E0 = eigsolve(H, copy(dv), 1, :SR; issymmetric=true)[1][1]

            df = lomc!(H, dv; post_step, laststep=3000).df

            # Shift estimate.
            Es, σs = mean_and_se(df.shift)
            s_low, s_high = Es - 2σs, Es + 2σs
            # Projected estimate.
            r = ratio_of_means(df.hproj, df.vproj)
            p_low, p_high = pquantile(r, [0.0015, 0.9985])

            @test s_low < E0 < s_high
            @test p_low < E0 < p_high
        end
    end

    @testset "Stochastic style comparison" begin
        add = BoseFS{5,5}((1,1,1,1,1))
        H = HubbardReal1D(add)
        E0 = -8.280991746582686

        seedCRNG!(1234)
        dv_st = DVec(add => 1; style=IsStochasticInteger())
        dv_th = DVec(add => 1; style=IsStochasticWithThreshold(1.0))
        dv_cx = DVec(add => 1 + im; style=IsStochastic2Pop())
        dv_dy = DVec(add => 1; style=IsDynamicSemistochastic())
        dv_de = DVec(add => 1; style=IsDeterministic())
        dv_dp = DVec(add => 1; style=IsDeterministic(ThresholdCompression()))
        dv_ex = DVec(add => 1; style=IsExplosive())

        dv_dn = DVec(add => 10; style=IsDynamicSemistochastic(compression=DoubleOrNothing()))
        dv_nr = DVec(add => 1; style=IsDynamicSemistochastic(spawning=WithoutReplacement()))
        dv_br = DVec(add => 1; style=IsDynamicSemistochastic(spawning=Bernoulli()))

        s_strat = DoubleLogUpdate(ζ=0.05, ξ=0.05^2/4, targetwalkers=100)
        s_strat_cx = DoubleLogUpdate(ζ=0.05, ξ=0.05^2/4, targetwalkers=100 + 100im)
        df_st = lomc!(H, dv_st; s_strat, laststep=2500).df
        df_th = lomc!(H, dv_th; s_strat, laststep=2500).df
        df_cx = lomc!(H, dv_cx; s_strat=s_strat_cx, laststep=2500).df
        df_dy = lomc!(H, dv_dy; s_strat, laststep=2500).df
        df_de = lomc!(H, dv_de; s_strat, laststep=2500).df
        df_dp = lomc!(H, dv_dp; s_strat, laststep=2500).df
        df_ex = lomc!(H, dv_ex; s_strat, laststep=2500).df

        df_dn = lomc!(H, dv_dn; s_strat, laststep=2500).df
        df_nr = lomc!(H, dv_nr; s_strat, laststep=2500).df
        df_br = lomc!(H, dv_br; s_strat, laststep=2500).df

        @test ("spawns", "deaths", "clones", "zombies", "annihilations") ⊆ names(df_st)
        @test ("spawns", "deaths", "clones", "zombies", "annihilations") ⊆ names(df_cx)
        @test "spawns" ∈ names(df_th)
        @test ("exact_steps", "inexact_steps", "spawns") ⊆ names(df_dy)
        @test "exact_steps" ∈ names(df_de)

        @test ("explosions", "ticks", "normal_steps",
               "explosive_spawns", "normal_spawns") ⊆ names(df_ex)
        @test ("exact_steps", "len_before") ⊆ names(df_dp)
        @test ("exact_steps", "len_before") ⊆ names(df_br)
        @test ("exact_steps", "len_before") ⊆ names(df_nr)
        @test ("exact_steps", "len_before") ⊆ names(df_dn)
        @test "len_before" ∉ names(df_st)
        @test "len_before" ∉ names(df_th)
        @test "len_before" ∉ names(df_cx)
        @test "len_before" ∉ names(df_de)

        E_st, σ_st = mean_and_se(df_st.shift[500:end])
        E_th, σ_th = mean_and_se(df_th.shift[500:end])
        E_cx, σ_cx = mean_and_se(real.(df_cx.shift[500:end]))
        E_dy, σ_dy = mean_and_se(df_dy.shift[500:end])
        E_de, σ_de = mean_and_se(df_de.shift[500:end])
        E_dp, σ_dp = mean_and_se(df_dp.shift[500:end])
        E_ex, σ_ex = mean_and_se(df_dp.shift[500:end])

        E_dn, σ_dn = mean_and_se(df_dn.shift[500:end])
        E_nr, σ_nr = mean_and_se(df_nr.shift[500:end])
        E_br, σ_br = mean_and_se(df_br.shift[500:end])

        # Stochastic noise depends on the method. Sampling without replacement makes a
        # small difference and is not consistently lower, so is not included here. A similar
        # thing happens with deterministic with compression and explosive spawns.
        @test σ_st > σ_th > σ_dn > σ_dy > σ_de
        # All estimates are fairly good.
        @test E_st ≈ E0 atol=3σ_st
        @test E_th ≈ E0 atol=3σ_th
        @test E_cx ≈ E0 atol=3σ_cx
        @test E_dy ≈ E0 atol=3σ_dy
        @test E_de ≈ E0 atol=3σ_de
        @test E_dp ≈ E0 atol=3σ_dp
        @test E_ex ≈ E0 atol=3σ_ex
        @test E_nr ≈ E0 atol=3σ_nr
        @test E_dn ≈ E0 atol=3σ_dn

        # For some reason, Bernoulli requires more walkers
        @test_broken E_br ≈ E0 atol=3σ_br
    end

    @testset "Initiator energies" begin
        add = BoseFS{10,10}((0,0,0,0,10,0,0,0,0,0))
        dv_no = DVec(
            add => 1;
            style=IsDynamicSemistochastic()
        )
        dv_i1 = InitiatorDVec(
            add => 1;
            initiator=Initiator(1),
            style=IsDynamicSemistochastic(),
        )
        dv_i2 = InitiatorDVec(
            add => 1;
            initiator=SimpleInitiator(1),
            style=IsDynamicSemistochastic(),
        )
        dv_i3 = InitiatorDVec(
            add => 1;
            initiator=CoherentInitiator(1),
            style=IsDynamicSemistochastic(),
        )

        @testset "Energies below the plateau & initiator bias" begin
            seedCRNG!(8008)

            H = HubbardMom1D(add; u=4.0)
            E0 = -9.251592973178997

            s_strat = DoubleLogUpdate(ζ=0.05, ξ=0.05^2/4, targetwalkers=300)
            laststep = 6_000
            dτ = 5e-4
            df_no = lomc!(H, copy(dv_no); s_strat, laststep, dτ).df
            df_i1 = lomc!(H, copy(dv_i1); s_strat, laststep, dτ).df
            df_i2 = lomc!(H, copy(dv_i2); s_strat, laststep, dτ).df
            df_i3 = lomc!(H, copy(dv_i3); s_strat, laststep, dτ).df

            E_no, σ_no = mean_and_se(df_no.shift[2000:end])
            E_i1, σ_i1 = mean_and_se(df_i1.shift[2000:end])
            E_i2, σ_i2 = mean_and_se(df_i2.shift[2000:end])
            E_i3, σ_i3 = mean_and_se(df_i3.shift[2000:end])

            # Garbage energy from no initiator.
            @test E_no < E0
            # Initiator has a bias.
            @test E_i1 > E0
            @test E_i2 > E0
            @test E_i3 > E0

            # Simple initiator has the largest bias.
            @test E_i2 > E_i1
            # Normal and coherent initiators are about the same.
            @test E_i1 ≈ E_i3 atol=max(3σ_i1, 3σ_i3)
        end

        @testset "Energies above the plateau" begin
            seedCRNG!(1337)

            H = HubbardMom1D(add)
            E0 = -16.36048582876015

            s_strat = DoubleLogUpdate(ζ=0.05, ξ=0.05^2/4, targetwalkers=3000)
            laststep = 2500
            dτ = 1e-2
            df_no = lomc!(H, copy(dv_no); s_strat, laststep, dτ).df
            df_i1 = lomc!(H, copy(dv_i1); s_strat, laststep, dτ).df
            df_i2 = lomc!(H, copy(dv_i2); s_strat, laststep, dτ).df
            df_i3 = lomc!(H, copy(dv_i3); s_strat, laststep, dτ).df

            E_no, σ_no = mean_and_se(df_no.shift[500:end])
            E_i1, σ_i1 = mean_and_se(df_i1.shift[500:end])
            E_i2, σ_i2 = mean_and_se(df_i2.shift[500:end])
            E_i3, σ_i3 = mean_and_se(df_i3.shift[500:end])

            # All estimates should be fairly good.
            @test E_no ≈ E0 atol=3σ_no
            @test E_i1 ≈ E0 atol=3σ_i1
            @test E_i2 ≈ E0 atol=3σ_i2
            @test E_i3 ≈ E0 atol=3σ_i3
        end
    end
end


text:
@parallel for i=1:100
    println(i)
end

max_stars_repo_name: Roger-luo/QCBM

max_issues_repo_name: Roger-luo/QCBM

max_forks_repo_name: Roger-luo/QCBM


text:
module TextUtils
import Base: *
using REPL
const LE = REPL.LineEdit
export is_linebreak, is_whitespace, is_word_char, TextChar, WordChar, WhitespaceChar, PunctuationChar, ObjectChar
export is_alphanumeric, is_alphabetic, is_uppercase, is_lowercase, is_punctuation
export is_object_end, is_object_start, is_whitespace_end, is_whitespace_start
export testbuf

"""
    Determine whether the buffer is currently at the end of a text object. Whitespace is not included as a text object.
"""
function is_object_start(buf)
    eof(buf) && return false
    start_pos = mark(buf)
    c1 = read(buf, Char) |> TextChar

    reset(buf)

    if c1 isa ObjectChar && start_pos == 0
        # beginning of line
        return true
    end

    mark(buf)

    skip(buf, -1)
    c2 = read(buf, Char) |> TextChar

    if c1 isa ObjectChar && typeof(c1) != typeof(c2)
        reset(buf)
        return true
    end
    return false
end

function is_whitespace_start(buf)
    eof(buf) && return false
    start_pos = mark(buf)
    c1 = read(buf, Char) |> TextChar

    reset(buf)

    if c1 isa WhitespaceChar && start_pos == 0
        # beginning of line
        return true
    end

    mark(buf)

    skip(buf, -1)
    c2 = read(buf, Char) |> TextChar

    if c1 isa WhitespaceChar && !(c2 isa WhitespaceChar)
        reset(buf)
        return true
    end
    return false
end

"""
    Whether the buffer is currently at the end of a text object. Whitespace is not included as a text object.
"""
function is_object_end(buf)
    eof(buf) && return false
    mark(buf)

    c1 = read(buf, Char) |> TextChar

    if eof(buf)
        reset(buf)
        return c1 isa ObjectChar
    end
    c2 = read(buf, Char) |> TextChar

    if c1 isa ObjectChar && typeof(c1) != typeof(c2)
        reset(buf)
        return true
    end
    reset(buf)
    return false
end

function is_whitespace_end(buf)
    eof(buf) && return false
    mark(buf)

    c1 = read(buf, Char) |> TextChar

    if eof(buf)
        reset(buf)
        return c1 isa WhitespaceChar
    end
    c2 = read(buf, Char) |> TextChar

    if c1 isa WhitespaceChar && !(c2 isa WhitespaceChar)
        reset(buf)
        return true
    end
    reset(buf)
    return false
end

"""
    Generate a buffer from s, but place its position where the pipe operator occurs in `s`
"""
function testbuf(s :: AbstractString) :: IOBuffer
    a, b = split(s, '|')
    buf = IOBuffer(a * b)
    seek(buf, length(a))
    return buf
end

"""
    A character as understood within the context of a vim object
"""
abstract type TextChar{T <: Char} end
# abstract type NonWordChar{T} <: TextChar{T} end
abstract type ObjectChar{T} <: TextChar{T} end

struct WhitespaceChar{T} <: TextChar{T}
    c :: T
end

struct WordChar{T} <: ObjectChar{T}
    c :: T
end

struct PunctuationChar{T} <: ObjectChar{T}
    c :: T
end

Base.convert(::Type{Char}, c::TextChar{Char}) = c.c
Base.promote_rule(::Type{<:TextChar}, ::Type{Char}) = Char

*(a::TextChar, b::AbstractChar) = *(promote(a, b)...)

function TextChar(c::T) where T <: Char
    return if is_whitespace(c)
        WhitespaceChar(c)
    elseif is_word_char(c)
        WordChar(c)
    elseif is_punctuation(c)
        PunctuationChar(c)
    else
        error("Char '$c' cannot be described by a TextChar")
    end
end

# Text helpers
is_linebreak(c::Char) = c in """\n"""
is_whitespace(c::Char) = c in """ \t\n"""
# non_word(c::Char) = LE.is_non_word_char(c)
is_word_char(c::Char) = is_alphanumeric(c) || c == '_'
# punct_char(c::Char) = LE.is_non_word_char(c) && !is_non_phrase_char(c)

function is_alphanumeric(c :: Char) :: Bool
    return c in ['a':'z';
                 'A':'Z';
                 '0':'9';]
end

function is_alphabetic(c :: Char) :: Bool
    return c in ['a':'z';
                 'A':'Z';]
end

function is_uppercase(c :: Char) :: Bool
    return c in ['A':'Z';]
end
function is_lowercase(c :: Char) :: Bool
    return c in ['a':'z';]
end

function is_punctuation(c :: Char) :: Bool
    return c in """`!@#\$%^&*()-_=+[]{}'\"/?\\|<>,.:;"""
end

end


text:
include("gridap_warmup.jl")

using PartitionedArrays
using PartitionedPoisson

#poisson(mode=:seq,nc=(10,10),np=(2,2))
#poisson(mode=:seq,nc=(10,10,10),np=(2,2,2))
poisson(mode=:mpi,nc=(10,10),np=(1,1),verbose=true)
poisson(mode=:mpi,nc=(10,10,10),np=(1,1,1),verbose=true)


text:
module RoundingIntegers

import Base: <, <=, +, -, *, ~, &, |, <<, >>, >>>, xor

export RSigned, RUnsigned, RInteger
export RInt8, RUInt8, RInt16, RUInt16, RInt32, RUInt32, RInt64, RUInt64,
       RInt128, RUInt128, RInt, RUInt

abstract type RSigned   <: Signed end
abstract type RUnsigned <: Unsigned end

const RInteger = Union{RSigned,RUnsigned}

primitive type RInt8 <: RSigned 8 end
primitive type RUInt8 <: RUnsigned 8 end
primitive type RInt16 <: RSigned 16 end
primitive type RUInt16 <: RUnsigned 16 end
primitive type RInt32 <: RSigned 32 end
primitive type RUInt32 <: RUnsigned 32 end
primitive type RInt64 <: RSigned 64 end
primitive type RUInt64 <: RUnsigned 64 end
primitive type RInt128 <: RSigned 128 end
primitive type RUInt128 <: RUnsigned 128 end

if Sys.WORD_SIZE == 32
    const RInt = RInt32
    const RUInt = RUInt32
else
    const RInt = RInt64
    const RUInt = RUInt64
end

itype(::Type{RInteger}) = Integer
itype(::Type{RSigned}) = Signed
itype(::Type{RUnsigned}) = Unsigned
itype(::Type{RInt8}) = Int8
itype(::Type{RUInt8}) = UInt8
itype(::Type{RInt16}) = Int16
itype(::Type{RUInt16}) = UInt16
itype(::Type{RInt32}) = Int32
itype(::Type{RUInt32}) = UInt32
itype(::Type{RInt64}) = Int64
itype(::Type{RUInt64}) = UInt64
itype(::Type{RInt128}) = Int128
itype(::Type{RUInt128}) = UInt128
itype(x::RInteger) = itype(typeof(x))

rtype(::Type{Signed}) = RSigned
rtype(::Type{Unsigned}) = RUnsigned
rtype(::Type{Int8}) = RInt8
rtype(::Type{UInt8}) = RUInt8
rtype(::Type{Int16}) = RInt16
rtype(::Type{UInt16}) = RUInt16
rtype(::Type{Int32}) = RInt32
rtype(::Type{UInt32}) = RUInt32
rtype(::Type{Int64}) = RInt64
rtype(::Type{UInt64}) = RUInt64
rtype(::Type{Int128}) = RInt128
rtype(::Type{UInt128}) = RUInt128
rtype(x::Union{Signed,Unsigned}) = rtype(typeof(x))

# RIntegers are largely about assignment; for usage, we convert to
# standard integers at the drop of a hat.

Base.promote_rule(::Type{RI}, ::Type{T}) where {T<:Number,RI<:RInteger} =
    promote_type(T, itype(RI))

Base.Signed(x::RSigned) = reinterpret(itype(x), x)
Base.Unsigned(x::RUnsigned) = reinterpret(itype(x), x)
RSigned(x::Signed) = reinterpret(rtype(x), x)
RUnsigned(x::Unsigned) = reinterpret(rtype(x), x)
Base.Integer(x::RSigned) = Signed(x)
Base.Integer(x::RUnsigned) = Unsigned(x)
RInteger(x::Signed) = RSigned(x)
RInteger(x::Unsigned) = RUnsigned(x)

Base.AbstractFloat(x::RInteger) = AbstractFloat(Integer(x))

RInt8(x::Int8) = reinterpret(RInt8, x)
RUInt8(x::UInt8) = reinterpret(RUInt8, x)
RInt16(x::Int16) = reinterpret(RInt16, x)
RUInt16(x::UInt16) = reinterpret(RUInt16, x)
RInt32(x::Int32) = reinterpret(RInt32, x)
RUInt32(x::UInt32) = reinterpret(RUInt32, x)
RInt64(x::Int64) = reinterpret(RInt64, x)
RUInt64(x::UInt64) = reinterpret(RUInt64, x)
RInt128(x::Int128) = reinterpret(RInt128, x)
RUInt128(x::UInt128) = reinterpret(RUInt128, x)

(::Type{R})(x::Integer) where R<:RInteger = R(convert(itype(R), x))
(::Type{R})(x::Float16) where R<:RInteger = R(round(itype(R), x))
(::Type{R})(x::BigFloat) where R<:RInteger = R(round(itype(R), x))
(::Type{R})(x::Rational) where R<:RInteger = R(round(itype(R), x))
(::Type{R})(x::Complex) where R<:RInteger = R(round(itype(R), x))
(::Type{R})(x::AbstractFloat) where R<:RInteger = R(round(itype(R), x))

@inline Base.convert(::Type{T}, x::RInteger) where {T<:Number} = convert(T, Integer(x))

# rem conversions
@inline Base.rem(x::T, ::Type{T}) where {T<:RInteger} = T
@inline Base.rem(x::Integer, ::Type{T}) where {T<:RInteger} = RInteger(rem(x, itype(T)))
# ambs
@inline Base.rem(x::BigInt, ::Type{T}) where {T<:RInteger} = error("no rounding BigInt available")

Base.flipsign(x::RSigned, y::RSigned) = RInteger(flipsign(Integer(x), Integer(y)))

<(x::RInteger, y::RInteger) = Integer(x) < Integer(y)
<=(x::RInteger, y::RInteger) = Integer(x) <= Integer(y)

Base.count_ones(x::RInteger) = count_ones(Integer(x))
Base.leading_zeros(x::RInteger) = leading_zeros(Integer(x))
Base.trailing_zeros(x::RInteger) = trailing_zeros(Integer(x))
Base.ndigits0z(x::RInteger) = Base.ndigits0z(Integer(x))
Base.ndigits0zpb(x::RSigned, b::Integer) = Base.ndigits0zpb(abs(Integer(x)), Int(b))
Base.ndigits0zpb(x::RUnsigned, b::Integer) = Base.ndigits0zpb(Integer(x), Int(b))

# A few operations preserve the type
-(x::RInteger) = RInteger(-Integer(x))
~(x::RInteger) = RInteger(~Integer(x))

>>(x::RInteger, y::Signed) = RInteger(Integer(x) >> y)
>>>(x::RInteger, y::Signed) = RInteger(Integer(x) >>> y)
<<(x::RInteger, y::Signed) = RInteger(Integer(x) << y)
>>(x::RInteger, y::Unsigned) = RInteger(Integer(x) >> y)
>>>(x::RInteger, y::Unsigned) = RInteger(Integer(x) >>> y)
<<(x::RInteger, y::Unsigned) = RInteger(Integer(x) << y)
# ambs
>>(x::RInteger, y::Int) = RInteger(Integer(x) >> y)
>>>(x::RInteger, y::Int) = RInteger(Integer(x) >>> y)
<<(x::RInteger, y::Int) = RInteger(Integer(x) << y)

+(x::T, y::T) where {T<:RInteger} = RInteger(Integer(x) + Integer(y))
-(x::T, y::T) where {T<:RInteger} = RInteger(Integer(x) - Integer(y))
*(x::T, y::T) where {T<:RInteger} = RInteger(Integer(x) * Integer(y))
(&)(x::T, y::T) where {T<:RInteger} = RInteger(Integer(x) & Integer(y))
(|)(x::T, y::T) where {T<:RInteger} = RInteger(Integer(x) | Integer(y))
xor(x::T, y::T) where {T<:RInteger} = RInteger(xor(Integer(x), Integer(y)))

Base.rem(x::T, y::T) where {T<:RInteger} = RInteger(rem(Integer(x), Integer(y)))
Base.mod(x::T, y::T) where {T<:RInteger} = RInteger(mod(Integer(x), Integer(y)))

Base.unsigned(x::RSigned) = RInteger(unsigned(Integer(x)))
Base.signed(x::RSigned)   = RInteger(signed(Integer(x)))

# traits
Base.typemin(::Type{T}) where {T<:RInteger} = RInteger(typemin(itype(T)))
Base.typemax(::Type{T}) where {T<:RInteger} = RInteger(typemax(itype(T)))
Base.widen(::Type{T}) where {T<:RInteger} = rtype(widen(itype(T)))

end # module


text:
# example to do...run with @btime macro.
using BenchmarkTools
function string_zeros(s::AbstractString)
    n = 1000_000
    x = s =="Int64" ? Vector{Int64}(undef,n) : Vector{Float64}(undef,n)
    return fill_zeros(x)
end
function fill_zeros(x)
    for i in 1:length(x)
        x[i]=0
    end
    return x
end
function arr_summary(x::Array{T}) where T <: Number
    r = zero(T)
    for i = 1:length(x)
        r = r + x[i] ^ 2
    end
    return r
end


text:
using LsqFit, CSV, Plots

# *** Fit Volumes to logistic model ***

@. model(age, pars) = pars[1] / (1 + exp(-pars[2] * (age - pars[3])))

obsVols = CSV.read(IOBuffer("""
age  vol
20   56
35   140
60   301
90   420
145  502
"""), delim=" ", ignorerepeated=true, copycols=true)

par0 = [600, 0.02, 60]

fit = curve_fit(model, obsVols.age, obsVols.vol, par0)

fit.param = [497.07, 0.05, 53.5]
fitX = 0:maximum(obsVols.age)*1.2
fitY = [fit.param[1] / (1 + exp(-fit.param[2] * (x - fit.param[3]))) for x in fitX]

plot(fitX, fitY, seriestype=:line, label = "Fitted values")
plot!(obsVols.age, obsVols.vol, seriestype = :scatter, label = "Obs values")
plot!(obsVols.age, fit.resid, seriestype=:bar, label = "Residuals")


text:
module Functional

using LinearAlgebra: Matrix, zero, similar
using LinearAlgebra, Printf
# using Roots
# using Optim
using Quadmath
# using ProfileView
# using InteractiveUtils

using ..Discrete

# const Float = Float64
# const Float = BigFloat
const Float = Float128

include("kernel.jl")
include("findmax.jl")

# using Plots
# function plotResidual(basis, proj, gmin, gmax, candidate=nothing, residual=nothing)
#     ω = LinRange(gmin, gmax, 1000)
#     y = [Residual(basis, proj, w) for w in ω]
#     p = plot(ω, y, xlims=(gmin, gmax))
#     if isnothing(candidate) == false
#         plot!(p, candidate, residual, seriestype=:scatter)
#     end
#     display(p)
#     readline()
# end

mutable struct Basis
    Λ::Float
    rtol::Float

    N::Int # number of basis
    grid::Vector{Float} # grid for the basis
    residual::Vector{Float} # achieved error by each basis
    Q::Matrix{Float} # K = Q*R
    proj::Matrix{Float} # the overlap of basis functions <K(g_i), K(g_j)>

    candidate::Vector{Float}
    candidateResidual::Vector{Float}

    function Basis(Λ, rtol)
        _Q = Matrix{Float}(undef, (0, 0))
        return new(Λ, rtol, 0, [], [], _Q, similar(_Q), [], [])
    end
end

function addBasis!(basis, proj, g0::Float)
    basis.N += 1
    if basis.N == 1
        idx = 1
        basis.grid = [g0,]
        basis.Q = zeros(Float, (basis.N, basis.N))
        basis.Q[1, 1] = 1 / sqrt(proj(basis.Λ, g0, g0))
        basis.proj = projKernel(basis, proj)
    else
        idxList = findall(x -> x > g0, basis.grid)
        # if ω is larger than any existing freqencies, then idx is an empty list
        idx = length(idxList) == 0 ? basis.N : idxList[1] # the index to insert the new frequency

        insert!(basis.grid, idx, g0)
        basis.proj = projKernel(basis, proj)
        _Q = copy(basis.Q)
        basis.Q = zeros(Float, (basis.N, basis.N))
        basis.Q[1:idx-1, 1:idx-1] = _Q[1:idx-1, 1:idx-1]
        basis.Q[1:idx-1, idx+1:end] = _Q[1:idx-1, idx:end]
        basis.Q[idx+1:end, 1:idx-1] = _Q[idx:end, 1:idx-1]
        basis.Q[idx+1:end, idx+1:end] = _Q[idx:end, idx:end]
        # println(maximum(abs.(GramSchmidt(basis, idx, g0) .- mGramSchmidt(basis, idx, g0))))
        basis.Q[idx, :] = mGramSchmidt(basis, idx, g0)
    end

    scanResidual!(basis, proj, g0, idx)
    insert!(basis.residual, idx, maximum(basis.candidateResidual)) # record error after the new grid is added
    return idx
end

function scanResidual!(basis, proj, g0, idx)
    grids = copy(basis.grid)
    if basis.grid[1] > Float(0)
        insert!(grids, 1, Float(0))
    end
    if basis.grid[end] < basis.Λ
        append!(grids, basis.Λ)
    end
    resize!(basis.candidate, length(grids) - 1)
    resize!(basis.candidateResidual, length(grids) - 1)
    # println(g0, " and ", idx)
    # println(grids)

    for i = 1:length(grids)-1 # because of the separation of scales, the grids far away from idx is rarely affected
        g = findCandidate(basis, proj, grids[i], grids[i+1])
        basis.candidate[i] = g
        basis.candidateResidual[i] = Residual(basis, proj, g)
    end
end

function printCandidate(basis, idx)
    lower = (idx == 1) ? 0 : basis.grid[idx-1]
    upper = (idx == basis.N) ? basis.Λ : basis.grid[idx+1]

    @printf("%3i : ω=%24.8f ∈ (%24.8f, %24.8f) -> error=%24.16g\n", basis.N, basis.grid[idx], lower, upper, basis.residual[idx])
end

function QR(Λ, rtol, proj, g0; N = nothing)
    basis = Basis(Λ, rtol)
    # println(g0)
    for g in g0
        idx = addBasis!(basis, proj, Float(g))
        # @printf("%3i : ω=%24.8f ∈ (%24.8f, %24.8f) -> error=%24.16g\n", 1, g, 0, Λ, basis.residual[idx])
        printCandidate(basis, idx)
    end

    # @code_warntype Residual(basis, proj, Float(1.0))
    # exit(0)
    maxResidual, ωi = findmax(basis.candidateResidual)
    # plotResidual(basis, proj, Float(0), Float(100), basis.candidate, basis.candidateResidual)

    while isnothing(N) ? maxResidual > rtol / 10 : basis.N < N

        newω = basis.candidate[ωi]
        idx = addBasis!(basis, proj, newω)
        printCandidate(basis, idx)
        # println(length(basis.grid))
        # println(idx)
        # lower = (idx == 1) ? 0 : basis.grid[idx - 1]
        # upper = (idx == basis.N) ? Λ : basis.grid[idx + 1]

        # @printf("%3i : ω=%24.8f ∈ (%24.8f, %24.8f) -> error=%24.16g\n", basis.N, newω, lower, upper, basis.residual[idx])
        # println("$(length(freq)) basis: ω=$(Float64(newω)) between ($(Float64(freq[idx - 1])), $(Float64(freq[idx + 1])))")
        # plotResidual(basis, proj, Float(0), Float(100), candidate, residual)
        maxResidual, ωi = findmax(basis.candidateResidual)
    end
    testOrthgonal(basis)
    # @printf("residual = %.10e, Fnorm/F0 = %.10e\n", residual, residualF(freq, Q, Λ))
    @printf("residual = %.10e\n", maximum(basis.candidateResidual))
    # plotResidual(basis, proj, Float(0), Float(100), basis.candidate, basis.candidateResidual)
    return basis
end

"""
q1=sum_j c_j K_j
q2=sum_k d_k K_k
return <q1, q2> = sum_jk c_j*d_k <K_j, K_k>
"""
projqq(basis, q1::Vector{Float}, q2::Vector{Float}) = q1' * basis.proj * q2

"""
<K(g_i), K(g_j)>
"""
function projKernel(basis, proj)
    K = zeros(Float, (basis.N, basis.N))
    for i = 1:basis.N
        for j = 1:basis.N
            K[i, j] = proj(basis.Λ, basis.grid[i], basis.grid[j])
        end
    end
    return K
end

"""
modified Gram-Schmidt process
"""
function mGramSchmidt(basis, idx, g::Float)
    qnew = zeros(Float, basis.N)
    qnew[idx] = 1

    for qi = 1:basis.N
        if qi == idx
            continue
        end
        q = basis.Q[qi, :]
        qnew -= projqq(basis, q, qnew) .* q  # <q, qnew> q
    end
    return qnew / sqrt(projqq(basis, qnew, qnew))
end

# """
# Gram-Schmidt process
# """
# function GramSchmidt(basis, idx, g::Float)
#     q0 = zeros(Float, basis.N)
#     q0[idx] = 1
#     qnew = copy(q0)

#     for qi in 1:basis.N
#         if qi == idx
#     continue
#     end
#         q = basis.Q[qi, :]
#         qnew -=  projqq(basis, q, q0) .* q
#     end

#     norm = sqrt(projqq(basis, qnew, qnew))
#     return qnew / norm
# end

function Residual(basis, proj, g::Float)
    # norm2 = proj(g, g) - \sum_i (<qi, K_g>)^2
    # qi=\sum_j Q_ij K_j ==> (<qi, K_g>)^2 = (\sum_j Q_ij <K_j, K_g>)^2 = \sum_jk Q_ij*Q_ik <K_j, K_g>*<K_k, Kg>

    KK = [proj(basis.Λ, gj, g) for gj in basis.grid]
    norm2 = proj(basis.Λ, g, g) - (norm(basis.Q * KK))^2

    # norm2 = proj(basis.Λ, g, g)
    # for j in 1:basis.N
    #     norm2 -= basisQ[j, :]
    # end
    return norm2 < 0 ? Float(0) : sqrt(norm2)
end

function testOrthgonal(basis)
    println("testing orthognalization...")
    II = basis.Q * basis.proj * basis.Q'
    maxerr = maximum(abs.(II - I))
    println("Max Orthognalization Error: ", maxerr)
end

"""
function build(dlrGrid, print::Bool = true)
    Construct discrete Lehmann representation

#Arguments:
- `dlrGrid`: struct that contains the information to construct the DLR grid. The following entries are required:
   Λ: the dimensionless scale β*Euv, rtol: the required relative accuracy, isFermi: fermionic or bosonic, symmetry: particle-hole symmetry/antisymmetry or none
- `print`: print the internal information or not
"""
function build(dlrGrid, print::Bool = true)
    print && println("Using the functional algorithm to build DLR ...")
    Λ = Float(dlrGrid.Λ)
    rtol = Float(dlrGrid.rtol)
    symmetry = dlrGrid.symmetry
    if symmetry == :ph
        print && println("Building ω grid ... ")
        ωBasis = QR(Λ, rtol, projPH_ω, [Float(0), Float(Λ)])
        ωGrid = ωBasis.grid
        rank = ωBasis.N
    elseif symmetry == :pha
        print && println("Building ω grid ... ")
        ωBasis = QR(Λ, rtol, projPHA_ω, [Float(Λ),])
        ωGrid = ωBasis.grid
        rank = ωBasis.N
    else
        error("Functional algorithm for the symmetry $symmetry has not yet been implemented!")
        # elseif type == :fermi
        #     println("Building ω grid ... ")
        #     ωBasis = QR(Λ, rtol, projPH_ω, [Float(0), Float(Λ)])
        #     ωGrid = vcat(-ωBasis.grid[end:-1:2], ωBasis.grid)
        #     rank = length(ωGrid)
        #     println("rank: $rank")
        #     println("Building τ grid ... ")
        #     τBasis = tauGrid(ωGrid, rank, Λ, rtol, :fermi)
        #     # τBasis = QR(Λ / 2, rtol / 10, projPH_τ, Float(0), N=ωBasis.N)
        #     println("Building n grid ... ")
        #     nBasis = MatFreqGrid(ωGrid, rank, Λ, :fermi)
    end

    ωGrid = Float64.(ωGrid)
    degree = 128
    τ = Discrete.τChebyGrid(dlrGrid, degree, print)
    kernel = Discrete.preciseKernelT(dlrGrid, τ, ωGrid, print)
    Discrete.testInterpolation(dlrGrid, τ, ωGrid, kernel, print)

    τIndex = Discrete.τnQR(kernel, rank, print)
    τGrid = sort(τ.grid[τIndex])

    nFineGrid, nFermiKernel, nBoseKernel = Discrete.preciseKernelΩn(dlrGrid, ωGrid, print)

    nFermiIndex = Discrete.τnQR(nFermiKernel, rank, print)
    nFermiGrid = sort(nFineGrid[nFermiIndex])

    nBoseIndex = Discrete.τnQR(nBoseKernel, rank, print)
    nBoseGrid = sort(nFineGrid[nBoseIndex])

    # τGrid = τBasis / Λ
    # τGrid = τBasis
    # nGrid = nBasis
    ########### output  ############################
    # @printf("%5s  %32s  %32s  %8s\n", "index", "real freq", "tau", "ωn")
    # for r = 1:rank
    #     @printf("%5i  %32.17g  %32.17g  %16i\n", r, ωGrid[r], τGrid[r], nGrid[r])
    # end
    ########### output  ############################
    print && @printf("%5s  %32s  %32s  %11s  %11s\n", "index", "real freq", "tau", "fermi ωn", "bose ωn")
    for r = 1:rank
        print && @printf("%5i  %32.17g  %32.17g  %16i %16i\n", r, ωGrid[r], τGrid[r], nFermiGrid[r], nBoseGrid[r])
    end

    return ωGrid, τGrid, nFermiGrid, nBoseGrid
    # return Dict([(:ω, ωGrid), (:τ, τGrid), (:ωn, nGrid)])
end

end

if abspath(PROGRAM_FILE) == @__FILE__
    # freq, Q = findBasis(1.0e-3, Float(100))
    # basis = QR(100, 1e-3)
    Λ = 1e10
    # Λ = 100
    # @time ωBasis = QR(Λ, 1e-13, projPH_ω, [Float(0), Float(Λ)])
    @time ωBasis = Functional.QR(Λ, 1e-12, projPHA_ω, [Float(Λ),])
    # @time τBasis = QR(Λ / 2, 1e-11, projPHA_τ, Float(0), N=ωBasis.N)
    # nBasis = MatFreqGrid(ωBasis.grid, ωBasis.N, Λ, :acorr)

    # @time basis = QR(100, 1e-10)
    # readline()
    # basis = QR(100, 1e-3)

end