include("ham.jl")
using ForwardDiff, DiffResults


function transition_fwd(o::HamFlow, ϵ, x)  
    d = o.d
    z, ρ, u, _= MixFlow.one_fwd(o, ϵ, ref_coord, x[1:d], x[d+1:2d], x[end])
    return [z; ρ; u] 
end

function transition_bwd(o::HamFlow, ϵ, x)  
    d = o.d
    z, ρ, u, _= MixFlow.one_bwd(o, ϵ, inv_ref_coord, x[1:d], x[d+1:2d], x[end])
    return [z; ρ; u] 
end

function one_fwd_tran_matrix(o::HamFlow, ϵ, x, Diffrs)
    Diffrs = ForwardDiff.jacobian!(Diffrs, x -> transition_fwd(o, ϵ, x), x)
    return DiffResults.value(Diffrs), DiffResults.jacobian(Diffrs)
end

function one_bwd_tran_matrix(o::HamFlow, ϵ, x, Diffrs)
    Diffrs = ForwardDiff.jacobian!(Diffrs, x -> transition_bwd(o, ϵ, x), x)
    return DiffResults.value(Diffrs), DiffResults.jacobian(Diffrs)
end


# f(x) = x.^2
# x = rand(4)
# # y = randn(4)
# result = DiffResults.JacobianResult(x)
# result = ForwardDiff.jacobian!(result, f, x)
# DiffResults.jacobian(value)

function flow_fwd_matrices(o::HamFlow, ϵ::Vector{T}, x, n_ref::Int64) where T<:Real
    d = o.d
    rs = DiffResults.JacobianResult(x)
    M = []
    Xs = Matrix{eltype(x)}(undef, n_ref, 2d+1)
    for i in 1: n_ref-1
        x, J = one_fwd_tran_matrix(o, ϵ, x, rs)
        Xs[i,:] .= x
        push!(M, J)
    end
    return Xs, M
end


function flow_bwd_matrices(o::HamFlow, ϵ::Vector{T}, x, n_ref::Int64) where T<:Real
    d = o.d
    rs = DiffResults.JacobianResult(x)
    M = []
    Xs = Matrix{eltype(x)}(undef, n_ref, 2d+1)
    for i in 1: n_ref-1
        x, J = one_bwd_tran_matrix(o, ϵ, x, rs)
        Xs[i,:] .= x
        push!(M, J)
    end
    return Xs, M
end
