export correlation_3op_2t, correlation_3op_1t, correlation_2op_2t, correlation_2op_1t

function _check_correlation_time_list(tlist::AbstractVector)
    any(t -> t == 0, tlist) ||
        throw(ArgumentError("The time list for calculating correlation function must contain the element `0`"))
    all(>=(0), tlist) ||
        throw(ArgumentError("All the elements in the time list for calculating correlation function must be positive."))
    return nothing
end

@doc raw"""
    correlation_3op_2t(H::AbstractQuantumObject,
        ψ0::Union{Nothing,QuantumObject},
        tlist::AbstractVector,
        τlist::AbstractVector,
        c_ops::Union{Nothing,AbstractVector,Tuple},
        A::QuantumObject,
        B::QuantumObject,
        C::QuantumObject;
        kwargs...)

Returns the two-time correlation function of three operators ``\hat{A}``, ``\hat{B}`` and ``\hat{C}``: ``\left\langle \hat{A}(t) \hat{B}(t + \tau) \hat{C}(t) \right\rangle`` for a given initial state ``|\psi_0\rangle``.

If the initial state `ψ0` is given as `nothing`, then the [`steadystate`](@ref) will be used as the initial state. Note that this is only implemented if `H` is constant ([`QuantumObject`](@ref)).
"""
function correlation_3op_2t(
        H::AbstractQuantumObject{HOpType},
        ψ0::Union{Nothing, QuantumObject{StateOpType}},
        tlist::AbstractVector,
        τlist::AbstractVector,
        c_ops::Union{Nothing, AbstractVector, Tuple},
        A::QuantumObject{Operator},
        B::QuantumObject{Operator},
        C::QuantumObject{Operator};
        kwargs...,
    ) where {HOpType <: Union{Operator, SuperOperator}, StateOpType <: Union{Ket, Operator}}
    # check tlist and τlist
    _check_correlation_time_list(tlist)
    _check_correlation_time_list(τlist)

    L = liouvillian(H, c_ops)
    if ψ0 isa Nothing
        ψ0 = steadystate(L)
    end

    # check the dimensions of Operator A, B, C, and ψ0
    # the dimension validity of L and ψ0 will be checked in mesolve
    check_mul_dimensions(A, ψ0)
    check_dimensions(A, B, C)

    kwargs2 = merge((saveat = collect(tlist),), (; kwargs...))

    # generate OperatorKet first, so sol.states will store OperatorKet (skip redundant reshaping)
    # note that doesn't need to do the first mesolve if tlist == [0] (for 2op_1t and 3op_1t cases)
    ρ0_vec = mat2vec(ket2dm(ψ0))
    ρt_list = (tlist == [0]) ? [ρ0_vec] : mesolve(L, ρ0_vec, tlist; kwargs2...).states

    C_A = sprepost(C, A)
    corr = map((t, ρt) -> mesolve(L, C_A * ρt, τlist .+ t, e_ops = [B]; kwargs...).expect[1, :], tlist, ρt_list)

    # make the output correlation Matrix align with QuTiP
    # 1st dimension corresponds to tlist
    # 2nd dimension corresponds to τlist
    return reduce(vcat, transpose.(corr))
end

@doc raw"""
    correlation_3op_1t(H::AbstractQuantumObject,
        ψ0::Union{Nothing,QuantumObject},
        τlist::AbstractVector,
        c_ops::Union{Nothing,AbstractVector,Tuple},
        A::QuantumObject,
        B::QuantumObject,
        C::QuantumObject;
        kwargs...)

Returns the two-time correlation function (with only one time coordinate ``\tau``) of three operators ``\hat{A}``, ``\hat{B}`` and ``\hat{C}``: ``\left\langle \hat{A}(0) \hat{B}(\tau) \hat{C}(0) \right\rangle`` for a given initial state ``|\psi_0\rangle``.

If the initial state `ψ0` is given as `nothing`, then the [`steadystate`](@ref) will be used as the initial state. Note that this is only implemented if `H` is constant ([`QuantumObject`](@ref)).
"""
function correlation_3op_1t(
        H::AbstractQuantumObject{HOpType},
        ψ0::Union{Nothing, QuantumObject{StateOpType}},
        τlist::AbstractVector,
        c_ops::Union{Nothing, AbstractVector, Tuple},
        A::QuantumObject{Operator},
        B::QuantumObject{Operator},
        C::QuantumObject{Operator};
        kwargs...,
    ) where {HOpType <: Union{Operator, SuperOperator}, StateOpType <: Union{Ket, Operator}}
    corr = correlation_3op_2t(H, ψ0, [0], τlist, c_ops, A, B, C; kwargs...)

    return corr[1, :] # 1 means tlist[1] = 0
end

@doc raw"""
    correlation_2op_2t(H::AbstractQuantumObject,
        ψ0::Union{Nothing,QuantumObject},
        tlist::AbstractVector,
        τlist::AbstractVector,
        c_ops::Union{Nothing,AbstractVector,Tuple},
        A::QuantumObject,
        B::QuantumObject;
        reverse::Bool=false,
        kwargs...)

Returns the two-time correlation function of two operators ``\hat{A}`` and ``\hat{B}`` : ``\left\langle \hat{A}(t + \tau) \hat{B}(t) \right\rangle`` for a given initial state ``|\psi_0\rangle``.

If the initial state `ψ0` is given as `nothing`, then the [`steadystate`](@ref) will be used as the initial state. Note that this is only implemented if `H` is constant ([`QuantumObject`](@ref)).

When `reverse=true`, the correlation function is calculated as ``\left\langle \hat{A}(t) \hat{B}(t + \tau) \right\rangle``.
"""
function correlation_2op_2t(
        H::AbstractQuantumObject{HOpType},
        ψ0::Union{Nothing, QuantumObject{StateOpType}},
        tlist::AbstractVector,
        τlist::AbstractVector,
        c_ops::Union{Nothing, AbstractVector, Tuple},
        A::QuantumObject{Operator},
        B::QuantumObject{Operator};
        reverse::Bool = false,
        kwargs...,
    ) where {HOpType <: Union{Operator, SuperOperator}, StateOpType <: Union{Ket, Operator}}
    C = eye(get_size(A.dimensions)[1], type = Operator(), dims = A.dims) # same as qeye_like(A), use A instead of H (cause H might be SuperOperator)
    if reverse
        corr = correlation_3op_2t(H, ψ0, tlist, τlist, c_ops, A, B, C; kwargs...)
    else
        corr = correlation_3op_2t(H, ψ0, tlist, τlist, c_ops, C, A, B; kwargs...)
    end

    return corr
end

@doc raw"""
    correlation_2op_1t(H::AbstractQuantumObject,
        ψ0::Union{Nothing,QuantumObject},
        τlist::AbstractVector,
        c_ops::Union{Nothing,AbstractVector,Tuple},
        A::QuantumObject,
        B::QuantumObject;
        reverse::Bool=false,
        kwargs...)

Returns the two-time correlation function (with only one time coordinate ``\tau``) of two operators ``\hat{A}`` and ``\hat{B}`` : ``\left\langle \hat{A}(\tau) \hat{B}(0) \right\rangle`` for a given initial state ``|\psi_0\rangle``.

If the initial state `ψ0` is given as `nothing`, then the [`steadystate`](@ref) will be used as the initial state. Note that this is only implemented if `H` is constant ([`QuantumObject`](@ref)).

When `reverse=true`, the correlation function is calculated as ``\left\langle \hat{A}(0) \hat{B}(\tau) \right\rangle``.
"""
function correlation_2op_1t(
        H::AbstractQuantumObject{HOpType},
        ψ0::Union{Nothing, QuantumObject{StateOpType}},
        τlist::AbstractVector,
        c_ops::Union{Nothing, AbstractVector, Tuple},
        A::QuantumObject{Operator},
        B::QuantumObject{Operator};
        reverse::Bool = false,
        kwargs...,
    ) where {HOpType <: Union{Operator, SuperOperator}, StateOpType <: Union{Ket, Operator}}
    corr = correlation_2op_2t(H, ψ0, [0], τlist, c_ops, A, B; reverse = reverse, kwargs...)

    return corr[1, :] # 1 means tlist[1] = 0
end
