export adjtdvp1vec!, adaptiveadjtdvp1vec!

using ITensors: position!, permute
using ITensors.ITensorMPS: set_nsite!, linkdims, check_hascommoninds

# Version with a single initial state
"""
    adjtdvp1vec!(
        solver, operator::MPS, initialstate::MPS, PH,
        Δt::Number, tf::Number, meas_stride::Number, sites; kwargs...
    )

Compute the time evolution, generated by the GKSL superoperator encoded in `H`, of the
operator `operator` in the Heisenberg picture, periodically measuring its expectation
value on the state `initialstate`.

# Arguments
- `solver`: a function which takes three arguments `A`, `t`, `B` (and possibly other keyword
    arguments) where `t` is a time step, `B` an ITensor and `A` a linear operator on `B`,
    returning the time-evolved `B`.
- `operator::MPS`: an MPS representing the initial value of the operator.
- `initialstate::MPS` an MPS representing the initial state of the system.
- `PH`: a ProjMPO-like object representing the GKSL "superoperator".
- `Δt::Number`: the time step for the evolution.
- `tf::Number`: the end time of the simulation.
- `meas_stride::Number`: time between each measurement.
- `sites`: an array of ITensor sites on which the MPSs and MPOs above are defined.

# Keyword arguments
- `io_file`: output file of the measurements.
- `ranks_file`: output file for the bond dimensions of the operator's MPS.
- `times_file`: output file for the simulation time.

# Other keyword options, passed to `KrylovKit.exponentiate`
- `exp_tol::Real` -> `tol`
- `krylovdim::Int`
- `maxiter::Int`
"""
function adjtdvp1vec!(
    solver,
    operator::MPS,
    initialstate::MPS,
    PH,
    Δt::Number,
    tf::Number,
    meas_stride::Number,
    sites;
    kwargs...,
)
    return adjtdvp1vec!(
        solver, operator, [initialstate], PH, Δt, tf, meas_stride, sites; kwargs...
    )
end

# Version with a vector of initial states but an MPO Hamiltonian
"""
    adjtdvp1vec!(
        solver, operator::MPS, initialstates::Vector{MPS}, H::MPO,
        Δt::Number, tf::Number, meas_stride::Number, sites; kwargs...
    )

Compute the time evolution, generated by the GKSL superoperator encoded in `H`, of the
operator `operator` in the Heisenberg picture, periodically measuring its expectation
value on the state `initialstate`.

# Arguments
- `solver`: a function which takes three arguments `A`, `t`, `B` (and possibly other keyword
    arguments) where `t` is a time step, `B` an ITensor and `A` a linear operator on `B`,
    returning the time-evolved `B`.
- `operator::MPS`: an MPS representing the initial value of the operator.
- `initialstates::Vector{MPS}` a list of MPS representing initial states.
- `H::MPO`: the MPO representing the GKSL "superoperator".
- `Δt::Number`: the time step for the evolution.
- `tf::Number`: the end time of the simulation.
- `meas_stride::Number`: time between each measurement.
- `sites`: an array of ITensor sites on which the MPSs and MPOs above are defined.

# Keyword arguments
- `io_file`: output file of the measurements.
- `ranks_file`: output file for the bond dimensions of the operator's MPS.
- `times_file`: output file for the simulation time.

# Other keyword options, passed to `KrylovKit.exponentiate`
- `exp_tol::Real` -> `tol`
- `krylovdim::Int`
- `maxiter::Int`
"""
function adjtdvp1vec!(
    solver,
    operator::MPS,
    initialstates::Vector{MPS},
    H::MPO,
    Δt::Number,
    tf::Number,
    meas_stride::Number,
    sites;
    kwargs...,
)
    return adjtdvp1vec!(
        solver, operator, initialstates, ProjMPO(H), Δt, tf, meas_stride, sites; kwargs...
    )
end

# Complete version (a vector of initial states _and_ a ProjMPO Hamiltonian)
"""
    adjtdvp1vec!(
        solver,
        operator::MPS,
        initialstates::Vector{MPS},
        PH,
        Δt::Number, tf::Number, meas_stride::Number, sites; kwargs...
    )

Compute the time evolution, generated by the GKSL superoperator encoded in `H`, of the
operator `operator` in the Heisenberg picture, periodically measuring its expectation
value on each state in the vector `initialstates`.

# Arguments
- `solver`: a function which takes three arguments `A`, `t`, `B` (and possibly other keyword
    arguments) where `t` is a time step, `B` an ITensor and `A` a linear operator on `B`,
    returning the time-evolved `B`.
- `operator::MPS`: an MPS representing the initial value of the operator.
- `initialstates::Vector{MPS}` a list of MPS representing initial states.
- `PH`: a ProjMPO-like object representing the GKSL "superoperator".
- `Δt::Number`: the time step for the evolution.
- `tf::Number`: the end time of the simulation.
- `meas_stride::Number`: time between each measurement.
- `sites`: an array of ITensor sites on which the MPSs and MPOs above are defined.

# Keyword arguments
- `io_file`: output file of the measurements.
- `ranks_file`: output file for the bond dimensions of the operator's MPS.
- `times_file`: output file for the simulation time.

# Other keyword options, passed to `KrylovKit.exponentiate`
- `exp_tol::Real` -> `tol`
- `krylovdim::Int`
- `maxiter::Int`
"""
function adjtdvp1vec!(
    solver,
    operator::MPS,
    initialstates::Vector{MPS},
    PH,
    Δt::Number,
    tf::Number,
    meas_stride::Number,
    sites;
    kwargs...,
)
    nsteps = Int(tf / Δt)
    exp_tol = get(kwargs, :exp_tol, 1e-14)
    krylovdim = get(kwargs, :krylovdim, 30)
    maxiter = get(kwargs, :maxiter, 100)
    io_file = get(kwargs, :io_file, nothing)
    ranks_file = get(kwargs, :io_ranks, nothing)
    times_file = get(kwargs, :io_times, nothing)
    decomp = get(kwargs, :which_decomp, "qr")
    initialstatelabels = get(kwargs, :initialstatelabels, string.(1:length(initialstates)))

    if get(kwargs, :progress, true)
        pbar = Progress(nsteps; desc="Evolving operator... ")
    else
        pbar = nothing
    end

    # Vectorized equations of motion usually are not defined by an anti-Hermitian operator
    # such as -im H in Schrödinger's equation, so we do not bother here with "unitary" or
    # "imaginary-time" evolution types. We just have a generic equation of the form
    # v'(t) = L v(t).

    io_handle = open(io_file, "w")
    columnheaders = ["time"]
    if length(initialstates) == 1
        push!(columnheaders, "exp_val_real", "exp_val_imag")
    else
        for l in initialstatelabels
            push!(columnheaders, "exp_val_$(l)_real", "exp_val_$(l)_imag")
        end
    end
    println(io_handle, join(columnheaders, ","))

    ranks_handle = writeheaders_ranks(ranks_file, length(operator))
    times_handle = writeheaders_stime(times_file)

    N = length(operator)

    # Prepare for first iteration.
    orthogonalize!(operator, 1)
    set_nsite!(PH, 1)
    position!(PH, operator, 1)

    current_time = zero(Δt)
    prev_t = zero(Δt)
    for s in 1:nsteps
        stime = @elapsed begin
            # In TDVP1 only one site at a time is modified, so we iterate on the sites
            # of the operator MPS, not its bonds.
            for (site, ha) in sweepnext(N; ncenter=1)
                # sweepnext(N) is an iterable object that evaluates to tuples of the form
                # (bond, ha) where bond is the bond number and ha is the half-sweep number.
                # The kwarg ncenter determines the end and turning points of the loop: if
                # it equals 1, then we perform a sweep on each single site.
                sweepdir = (ha == 1 ? "right" : "left")
                tdvp_site_update!(
                    solver,
                    PH,
                    operator,
                    site,
                    0.5Δt;
                    current_time=(ha == 1 ? current_time + 0.5Δt : current_time + Δt),
                    sweepdir=sweepdir,
                    which_decomp=decomp,
                    hermitian=false,
                    exp_tol=exp_tol,
                    krylovdim=krylovdim,
                    maxiter=maxiter,
                )
            end
        end
        # The sweep has ended, so we update the current time.
        current_time += Δt

        !isnothing(pbar) && ProgressMeter.next!(
            pbar;
            showvalues=[
                ("t", current_time),
                ("Δt step time", round(stime; digits=3)),
                ("Max bond-dim", maxlinkdim(operator)),
            ],
        )

        # We actually want to measure (i.e. contract the MPS) not at each time step, but
        # on each k-th one where k = meas_stride/time_step.
        # We save the previous time at which we computed measurements in prev_t, and each
        # time a sweep and we check if current_time - prev_t = meas_stride.
        # If it is so, then we go on and compute the expectation value (and update prev_t).
        if (current_time - prev_t ≈ meas_stride || current_time == 0)
            data = [current_time]

            expvals = [inner(s, operator) for s in initialstates]
            for x in expvals
                push!(data, real(x), imag(x))
            end

            println(io_handle, join(data, ","))
            flush(io_handle)

            println(ranks_handle, join([current_time; linkdims(operator)], ","))
            flush(ranks_handle)

            printoutput_stime(times_handle, stime)

            prev_t = current_time
        end
    end  # of the time evolution.

    !isnothing(io_file) && close(io_handle)
    !isnothing(ranks_file) && close(ranks_handle)
    !isnothing(times_file) && close(times_handle)

    return nothing
end

"""
    adaptiveadjtdvp1vec!(
        solver,
        operator::MPS,
        initialstate::MPS,
        L::Vector{MPO},
        Δt::Number,
        tf::Number,
        meas_stride::Number,
        sites;
        kwargs...,
    )

Like `adjtdvp1vec!`, but grows the bond dimensions of the MPS along the time evolution until
a certain convergence criterium is met.

See [`adjtdvp1vec!`](@ref).
"""
function adaptiveadjtdvp1vec!(
    solver,
    operator::MPS,
    initialstate::MPS,
    Ls::Vector{MPO},
    Δt::Number,
    tf::Number,
    meas_stride::Number,
    sites;
    kwargs...,
)
    for L in Ls
        check_hascommoninds(siteinds, L, operator)
        check_hascommoninds(siteinds, L, operator')
    end
    Ls .= permute.(Ls, Ref((linkind, siteinds, linkind)))
    PLs = ProjMPOSum(Ls)
    return adaptiveadjtdvp1vec!(
        solver, operator, initialstate, PLs, Δt, tf, meas_stride, sites; kwargs...
    )
end

"""
    adaptiveadjtdvp1vec!(
        solver,
        operator::MPS,
        initialstate::MPS,
        L::MPO,
        Δt::Number,
        tf::Number,
        meas_stride::Number,
        sites;
        kwargs...,
    )

Like `adjtdvp1vec!`, but grows the bond dimensions of the MPS along the time evolution until
a certain convergence criterium is met.

See [`adjtdvp1vec!`](@ref).
"""
function adaptiveadjtdvp1vec!(
    solver,
    operator::MPS,
    initialstate::MPS,
    L::MPO,
    Δt::Number,
    tf::Number,
    meas_stride::Number,
    sites;
    kwargs...,
)
    return adaptiveadjtdvp1vec!(
        solver, operator, initialstate, ProjMPO(L), Δt, tf, meas_stride, sites; kwargs...
    )
end

"""
    adaptiveadjtdvp1vec!(
        solver,
        operator::MPS,
        initialstate::MPS,
        PH,
        Δt::Number,
        tf::Number,
        meas_stride::Number,
        sites;
        kwargs...,
    )

Like `adjtdvp1vec!`, but grows the bond dimensions of the MPS along the time evolution until
a certain convergence criterium is met.

See [`adjtdvp1vec!`](@ref).
"""
function adaptiveadjtdvp1vec!(
    solver,
    operator::MPS,
    initialstates::Vector{MPS},
    PH,
    Δt::Number,
    tf::Number,
    meas_stride::Number,
    sites;
    kwargs...,
)
    nsteps = Int(tf / Δt)
    exp_tol = get(kwargs, :exp_tol, 1e-14)
    krylovdim = get(kwargs, :krylovdim, 30)
    maxiter = get(kwargs, :maxiter, 100)
    io_file = get(kwargs, :io_file, nothing)
    ranks_file = get(kwargs, :io_ranks, nothing)
    times_file = get(kwargs, :io_times, nothing)
    convergence_factor_bonddims = get(kwargs, :convergence_factor_bonddims, 1e-4)
    max_bond = get(kwargs, :max_bond, maxlinkdim(operator))
    decomp = get(kwargs, :which_decomp, "qr")
    initialstatelabels = get(kwargs, :initialstatelabels, string.(1:length(initialstates)))

    if get(kwargs, :progress, true)
        pbar = Progress(nsteps; desc="Evolving operator... ")
    else
        pbar = nothing
    end

    # Vectorized equations of motion usually are not defined by an anti-Hermitian operator
    # such as -im H in Schrödinger's equation, so we do not bother here with "unitary" or
    # "imaginary-time" evolution types. We just have a generic equation of the form
    # v'(t) = L v(t).

    io_handle = open(io_file, "w")
    columnheaders = ["time"]
    if length(initialstates) == 1
        push!(columnheaders, "exp_val_real", "exp_val_imag")
    else
        for l in initialstatelabels
            push!(columnheaders, "exp_val_$(l)_real", "exp_val_$(l)_imag")
        end
    end
    println(io_handle, join(columnheaders, ","))

    ranks_handle = writeheaders_ranks(ranks_file, length(operator))
    times_handle = writeheaders_stime(times_file)

    N = length(operator)

    current_time = zero(Δt)
    prev_t = zero(Δt)
    for s in 1:nsteps
        orthogonalize!(operator, 1)
        set_nsite!(PH, 1)
        position!(PH, operator, 1)

        @debug "[Step $s] Attempting to grow the bond dimensions."
        adaptbonddimensions!(operator, PH, max_bond, convergence_factor_bonddims)

        stime = @elapsed begin
            # In TDVP1 only one site at a time is modified, so we iterate on the sites
            # of the operator MPS, not its bonds.
            for (site, ha) in sweepnext(N; ncenter=1)
                # sweepnext(N) is an iterable object that evaluates to tuples of the form
                # (bond, ha) where bond is the bond number and ha is the half-sweep number.
                # The kwarg ncenter determines the end and turning points of the loop: if
                # it equals 1, then we perform a sweep on each single site.
                sweepdir = (ha == 1 ? "right" : "left")
                tdvp_site_update!(
                    solver,
                    PH,
                    operator,
                    site,
                    0.5Δt;
                    current_time=(ha == 1 ? current_time + 0.5Δt : current_time + Δt),
                    sweepdir=sweepdir,
                    which_decomp=decomp,
                    hermitian=false,
                    exp_tol=exp_tol,
                    krylovdim=krylovdim,
                    maxiter=maxiter,
                )
            end
        end
        current_time += Δt

        !isnothing(pbar) && ProgressMeter.next!(
            pbar;
            showvalues=[
                ("t", current_time),
                ("Δt step time", round(stime; digits=3)),
                ("Max bond-dim", maxlinkdim(operator)),
            ],
        )

        # We actually want to measure (i.e. contract the MPS) not at each time step, but
        # on each k-th one where k = meas_stride/time_step.
        # We save the previous time at which we computed measurements in prev_t, and each
        # time a sweep and we check if current_time - prev_t = meas_stride.
        # If it is so, then we go on and compute the expectation value (and update prev_t).
        if (current_time - prev_t ≈ meas_stride || current_time == 0)
            data = [current_time]

            expvals = [inner(s, operator) for s in initialstates]
            for x in expvals
                push!(data, real(x), imag(x))
            end

            println(io_handle, join(data, ","))
            flush(io_handle)

            println(ranks_handle, join([current_time; linkdims(operator)], ","))
            flush(ranks_handle)

            printoutput_stime(times_handle, stime)

            prev_t = current_time
        end
    end  # of the time evolution.

    !isnothing(io_file) && close(io_handle)
    !isnothing(ranks_file) && close(ranks_handle)
    !isnothing(times_file) && close(times_handle)

    return nothing
end
