export adaptbonddimensions!

using ITensors: position!
using ITensors.ITensorMPS: AbstractProjMPO, set_nsite!, setleftlim!, setrightlim!

"""
    bondconvergencemeasure(PH::AbstractProjMPO, state::MPS, bond::Integer)

Return a measure of the convergence for the bond dimension on the bond (`bond`, `bond+1`)
of the MPS `state` for a time-evolution determined by `PH`. See [1] for details.

# References:
[1] Dunnett, Angus J. and Chin, Alex W. (2020).
“Dynamically Evolving Bond-Dimensions within the one-site
Time-Dependent-Variational-Principle method for Matrix Product States: Towards efficient
simulation of non-equilibrium open quantum dynamics”
https://doi.org/10.48550/arXiv.2007.13528
"""
function bondconvergencemeasure(PH::AbstractProjMPO, v::MPS, bond::Integer)::Real
    #orthogonalize!(v, bond) # this causes trouble

    set_nsite!(PH, 1)
    position!(PH, v, bond)
    H1 = PH(v[bond])

    Q, R = factorize(
        v[bond], uniqueinds(v[bond], v[bond + 1]); ortho="left", which_decomp="qr"
    )
    vv = copy(v)
    vv[bond] = Q
    vv[bond + 1] *= R
    setleftlim!(vv, bond)
    setrightlim!(vv, bond + 1)

    position!(PH, vv, bond + 1)
    H2 = PH(vv[bond + 1])

    set_nsite!(PH, 0)
    position!(PH, vv, bond + 1) # Force recalculation
    K = PH(R)

    return norm(H1)^2 + norm(H2)^2 + norm(K)^2
end

"""
    function adaptbonddimensions!(
        v::MPS, PH::AbstractProjMPO, max_bond::Int, convergence_factor_bonddims::Real
    )

If necessary, enlarge the bond dimensions of the MPS `v`, so that it can dynamically and
inhomogeneously restructure itself as the complexity of the correlations grow during the
time evolution generated by `H` [1].

# References:
[1] Dunnett, Angus J. and Chin, Alex W. (2020).
“Dynamically Evolving Bond-Dimensions within the one-site
Time-Dependent-Variational-Principle method for Matrix Product States: Towards efficient
simulation of non-equilibrium open quantum dynamics”
https://doi.org/10.48550/arXiv.2007.13528
"""
function adaptbonddimensions!(
    v::MPS, PH::AbstractProjMPO, max_bond::Int, convergence_factor_bonddims::Real
)
    maxbonddimensions = maxlinkdims(v, max_bond)
    for bond in 1:(length(v) - 1)
        if linkdim(v, bond) < maxbonddimensions[bond]
            # Skip all this if the bond is already at (or above!) the maximum
            # allowed value.
            f = bondconvergencemeasure(PH, v, bond)
            while true # (do-while block emulation)
                # Increase the bond dimension by 1, check convergence, repeat if needed.
                vcopy = copy(v)
                new_bonddim = growbond!(vcopy, bond)
                # Update PH using the new sites at `bond` and `bond + 1`:
                recompute!(PH, vcopy, bond)
                new_f = bondconvergencemeasure(PH, vcopy, bond)

                # If new_f / f ≈ 1, within the given threshold, then f was already OK
                # and we discard the new MPS, keeping the non-enlarged state.
                if (new_f / f - 1 > convergence_factor_bonddims && new_bonddim ≤ max_bond)
                    d = growbond!(v, bond) # Actually increase v's bond dimensions
                    #d = linkdim(v, bond)
                    @debug "[Bond ($bond,$(bond+1))] f($d)/f($(d-1)) - 1 > " *
                        "$convergence_factor_bonddims: increasing dimension to $d."
                    f = new_f # and begin a new iteration of the while block.
                else
                    # Bring PH back to its form before v's bond dimension was increased.
                    recompute!(PH, v, bond)
                    @debug "[Bond ($bond,$(bond+1))] Convergence reached at d=$(linkdim(v, bond))."
                    break # and proceed with the next bond in the MPS.
                end
            end
        else
            @debug "[Bond ($bond,$(bond+1))] Max dimension reached. Skipping."
        end
    end
    return nothing
end
