using OrdinaryDiffEq, Test, LinearAlgebra

T = 1000.0; Ttr = 0.0;
d0=1e-9; threshold=10^4*d0;
dt = 0.1; diff_eq_kwargs = Dict()

@inbounds function eom_lorenz!(du, u, p, t)
  σ = 10.0; ρ = 28.0; β = 8/3
  du[1] = σ*(u[2]-u[1])
  du[2] = u[1]*(ρ-u[3]) - u[2]
  du[3] = u[1]*u[2] - β*u[3]
end

prob = OrdinaryDiffEq.ODEProblem(eom_lorenz!, [0., 10.0, 0], (zero(T), T))
integ1 = init(prob, Tsit5(); diff_eq_kwargs...)
@. prob.u0 = prob.u0 + d0/sqrt(3)
integ2 = init(prob, Tsit5(); diff_eq_kwargs...)
integ1.opts.advance_to_tstop=true
integ2.opts.advance_to_tstop=true

dist = d0
λ = zero(eltype(integ1.u))
i = 0
tvector = dt:dt:T

for τ in tvector
  # evolve until rescaling:
  push!(integ1.opts.tstops, τ); step!(integ1)
  push!(integ2.opts.tstops, τ); step!(integ2)
  dist = norm(integ1.u .- integ2.u)
  # Rescale:
  if dist ≥ threshold
    # Rescale and reset everything:
    integ2.u .= integ1.u #.+ (integ2.u .- integ1.u)./a
    u_modified!(integ2, true)
    set_proposed_dt!(integ2,integ1)
    break
  end
end

τ = tvector[end]
push!(integ1.opts.tstops, τ); step!(integ1)
push!(integ2.opts.tstops, τ); step!(integ2)
@test dist = norm(integ1.u .- integ2.u) == 0
