include("mf_setting.jl")
include("../../inference/shadowing/window.jl")
using JLD
using TickTock

###############################
# compute delta 
###############################
Random.seed!(1)
nsample = 50
n_ref = 650
E_fwd = zeros(nsample, n_ref - 1)
E_bwd = zeros(nsample, n_ref - 1)
prog_bar = ProgressMeter.Progress(
    nsample; dt=0.5, barglyphs=ProgressMeter.BarGlyphs("[=> ]"), barlen=50, color=:yellow
)
@threads for i in 1:nsample
    ef, eb = single_err_traj(o, a, n_ref)
    E_fwd[i, :] .= ef
    E_bwd[i, :] .= eb
    ProgressMeter.next!(prog_bar)
end

JLD.save("result/delta.jld", "E_fwd", E_fwd, "E_bwd", E_bwd)

###############################
# compute epsilon
###############################
Random.seed!(1)
nsample = 10

δ = 1e-14
n_ref = 651
Ns = [10, 20, 50, 100, 200, 400, 650]

W_fwd = zeros(nsample, length(Ns))
W_bwd = zeros(nsample, length(Ns))
T_fwd = zeros(nsample, length(Ns))
T_bwd = zeros(nsample, length(Ns))

prog_bar = ProgressMeter.Progress(
    nsample; dt=0.5, barglyphs=ProgressMeter.BarGlyphs("[=> ]"), barlen=50, color=:yellow
)
for i in 1:nsample
    z0 = o.q_sampler(d) .* D .+ μ
    ρ0 = randn(d)
    u0 = rand()
    x0 = [z0; ρ0; u0]

    Xs_fwd, M_fwd = MixFlow.flow_fwd_matrices(o, a.leapfrog_stepsize, x0, n_ref)
    Xs_bwd, M_bwd = MixFlow.flow_bwd_matrices(o, a.leapfrog_stepsize, x0, n_ref)

    count = 0
    for j in eachindex(Ns)
        count += 1
        M1 = construct_shadow_matrix(M_fwd[1:(Ns[j] - 1)])
        tick()
        ϵ1 = shadowing_window(M1, δ)
        t1 = tok()
        T_fwd[i, count] = t1

        M2 = construct_shadow_matrix(M_bwd[1:(Ns[j] - 1)])
        tick()
        ϵ2 = shadowing_window(M2, δ)
        t2 = tok()
        T_bwd[i, count] = t2

        W_fwd[i, j] = ϵ1
        W_bwd[i, j] = ϵ2
    end
    # update progress bar
    ProgressMeter.next!(prog_bar)
end

JLD.save(
    "result/windows.jld",
    "W_fwd",
    W_fwd,
    "W_bwd",
    W_bwd,
    "Ns",
    Ns,
    "T_fwd",
    T_fwd,
    "T_bwd",
    T_bwd,
)
