# Use harmonic oscillator system for both sesolve and mesolve
const N_ad = 20
const a_ad = destroy(N_ad)
const ψ0_ad = fock(N_ad, 0)
const tlist_ad = range(0, 40, 100)
const ad_a = a_ad' * a_ad

# ---- SESOLVE ----
# For direct Forward differentiation
function my_f_sesolve_direct(p)
    H = p[1] * a_ad' * a_ad + p[2] * (a_ad + a_ad')
    sol = sesolve(H, ψ0_ad, tlist_ad, progress_bar = Val(false))
    return real(expect(ad_a, sol.states[end]))
end

# For SciMLSensitivity.jl (reverse mode with Mooncake and Enzyme)
coef_Δ(p, t) = p[1]
coef_F(p, t) = p[2]
const H_evo = QobjEvo(a_ad' * a_ad, coef_Δ) + QobjEvo(a_ad + a_ad', coef_F)

function my_f_sesolve(p, sensealg)
    sol = sesolve(
        H_evo,
        ψ0_ad,
        tlist_ad,
        progress_bar = Val(false),
        params = p,
        sensealg = sensealg,
    )
    return real(expect(ad_a, sol.states[end]))
end
const my_f_sesolve_bsa_enzyme = Base.Fix{2}(my_f_sesolve, BacksolveAdjoint(autojacvec = EnzymeVJP()))
const my_f_sesolve_bsa_mooncake = Base.Fix{2}(my_f_sesolve, BacksolveAdjoint(autojacvec = MooncakeVJP()))

# ---- MESOLVE ----
# For direct Forward differentiation
function my_f_mesolve_direct(p)
    H = p[1] * a_ad' * a_ad + p[2] * (a_ad + a_ad')
    c_ops = [sqrt(p[3]) * a_ad]
    sol = mesolve(H, ψ0_ad, tlist_ad, c_ops, progress_bar = Val(false))
    return real(expect(ad_a, sol.states[end]))
end

# For SciMLSensitivity.jl (reverse mode with Mooncake and Enzyme)
coef_γ(p, t) = sqrt(p[3])
const c_ops_ad = [QobjEvo(a_ad, coef_γ)]
const L_ad = liouvillian(H_evo, c_ops_ad)

function my_f_mesolve(p, sensealg)
    sol = mesolve(
        L_ad,
        ψ0_ad,
        tlist_ad,
        progress_bar = Val(false),
        params = p,
        sensealg = sensealg,
    )
    return real(expect(ad_a, sol.states[end]))
end
const my_f_mesolve_bsa_enzyme = Base.Fix{2}(my_f_mesolve, BacksolveAdjoint(autojacvec = EnzymeVJP()))
const my_f_mesolve_bsa_mooncake = Base.Fix{2}(my_f_mesolve, BacksolveAdjoint(autojacvec = MooncakeVJP()))

# Parameters for benchmarks
const params_sesolve = [1.0, 1.0]
const params_mesolve = [1.0, 1.0, 1.0]

function benchmark_autodiff!(SUITE)
    # Benchmark sesolve - Forward
    SUITE["Autodiff"]["sesolve"]["Forward"] = @benchmarkable ForwardDiff.gradient($my_f_sesolve_direct, $params_sesolve)

    # Benchmark sesolve - Reverse (Enzyme)
    SUITE["Autodiff"]["sesolve"]["Reverse (Enzyme)"] = @benchmarkable Enzyme.autodiff(
        Enzyme.set_runtime_activity(Enzyme.Reverse),
        Const($my_f_sesolve_bsa_enzyme),
        Active,
        Duplicated($params_sesolve, dparams_sesolve),
    ) setup = (dparams_sesolve = Enzyme.make_zero($params_sesolve))

    # Benchmark sesolve - Reverse (Mooncake)
    cache_sesolve_mooncake = Mooncake.prepare_gradient_cache(my_f_sesolve_bsa_mooncake, params_sesolve)
    SUITE["Autodiff"]["sesolve"]["Reverse (Mooncake)"] = @benchmarkable Mooncake.value_and_gradient!!(grad_cache1, $my_f_sesolve_bsa_mooncake, $params_sesolve) setup = (grad_cache1 = $cache_sesolve_mooncake)

    # Benchmark mesolve - Forward
    SUITE["Autodiff"]["mesolve"]["Forward"] = @benchmarkable ForwardDiff.gradient($my_f_mesolve_direct, $params_mesolve)

    # Benchmark mesolve - Reverse (Enzyme)
    SUITE["Autodiff"]["mesolve"]["Reverse (Enzyme)"] = @benchmarkable Enzyme.autodiff(
        Enzyme.set_runtime_activity(Enzyme.Reverse),
        Const($my_f_mesolve_bsa_enzyme),
        Active,
        Duplicated($params_mesolve, dparams_mesolve),
    ) setup = (dparams_mesolve = Enzyme.make_zero($params_mesolve))

    # Benchmark mesolve - Reverse (Mooncake)
    cache_mesolve_mooncake = Mooncake.prepare_gradient_cache(my_f_mesolve_bsa_mooncake, params_mesolve)
    SUITE["Autodiff"]["mesolve"]["Reverse (Mooncake)"] = @benchmarkable Mooncake.value_and_gradient!!(grad_cache2, $my_f_mesolve_bsa_mooncake, $params_mesolve) setup = (grad_cache2 = $cache_mesolve_mooncake)

    return nothing
end
