# =============================================================================
# Hairer-Wanner initial timestep estimation.
# Internal functions _ode_initdt_iip/_ode_initdt_oop implement the algorithm.
# When g !== nothing, stochastic diffusion terms are folded into the estimate:
#   d₁ uses max(|f₀±3g₀|)/sk instead of f₀/sk
#   d₂ uses max(|Δf±ΔgMax|)/sk instead of Δf/sk
# =============================================================================

# Coerce `a == b` to a scalar `Bool`. Some array wrappers (notably PyCall
# `PyObject` of arrays — JuliaPy/PyCall.jl#900) return `Vector{Bool}` from `==`,
# which is not valid in a boolean context. See OrdinaryDiffEq.jl#1402.
@inline function _bool_equal(a, b)
    r = a == b
    return r isa Bool ? r : all(r)
end

@muladd function _ode_initdt_iip(
        u0, t, tdir, dtmax, abstol, reltol, internalnorm,
        prob, g, noise_prototype, order, integrator
    )
    _tType = eltype(t)
    f = prob.f
    p = integrator.p
    oneunit_tType = oneunit(t)
    dtmax_tdir = tdir * dtmax

    dtmin = nextfloat(max(integrator.opts.dtmin, convert(_tType, oneunit_tType * eps(SciMLBase.value(t)))))
    smalldt = max(dtmin, convert(_tType, oneunit_tType * 1 // 10^(6)))

    if integrator.isdae
        result_dt = tdir * max(smalldt, dtmin)
        @SciMLMessage(
            lazy"Using default small timestep for DAE: dt = $(result_dt)",
            integrator.opts.verbose, :shampine_dt
        )
        return result_dt
    end

    # With zero continuous states there is nothing to estimate from (and indexing
    # into the empty state would throw), so fall back to a finite default dt.
    if isempty(u0)
        result_dt = tdir * max(smalldt, dtmin)
        @SciMLMessage(
            lazy"Empty initial state, using default small timestep: dt = $(result_dt)",
            integrator.opts.verbose, :shampine_dt
        )
        return result_dt
    end

    if eltype(u0) <: Number && !(integrator.alg isa CompositeAlgorithm)
        cache = get_tmp_cache(integrator)
        sk = first(cache)
        if u0 isa Array && abstol isa Number && reltol isa Number
            @inbounds @simd ivdep for i in eachindex(u0)
                sk[i] = abstol + internalnorm(u0[i], t) * reltol
            end
        else
            @.. broadcast = false sk = abstol + internalnorm(u0, t) * reltol
        end
    else
        if u0 isa Array && abstol isa Number && reltol isa Number
            sk = similar(u0, typeof(internalnorm(first(u0), t) * reltol))
            @inbounds @simd ivdep for i in eachindex(u0)
                sk[i] = abstol + internalnorm(u0[i], t) * reltol
            end
        else
            sk = @.. broadcast = false abstol + internalnorm(u0, t) * reltol
        end
    end

    if get_current_isfsal(integrator.alg, integrator.cache) &&
            integrator isa ODEIntegrator
        # Right now DelayDiffEq has issues with fsallast not being initialized
        f₀ = integrator.fsallast
        f(f₀, u0, p, t)
    else
        # TODO: use more caches
        # When time is unitless, `f` writes `du` with `eltype(u0)` and the
        # FunctionWrapper signature emitted by `promote_f` reflects that. Dividing
        # by `oneunit_tType` is only meaningful for unit-aware time (e.g. Unitful);
        # for plain numeric `t` it would spuriously promote (e.g. Float32 / Float64
        # -> Float64) and produce a `du` whose eltype no longer matches the wrapper.
        if recursive_unitless_eltype(u0) === eltype(u0)
            f₀ = zero(u0)
        elseif u0 isa Array && eltype(u0) isa Number
            T = eltype(first(u0) / oneunit_tType)
            f₀ = similar(u0, T)
            fill!(f₀, zero(T))
        else
            f₀ = zero.(u0 ./ oneunit_tType)
        end
        f(f₀, u0, p, t)
    end

    # TODO: use more caches
    #tmp = cache[2]

    if u0 isa Array
        if length(u0) > 0
            T = typeof(u0[1] / sk[1])
        else
            T = promote_type(eltype(u0), eltype(sk))
        end
        tmp = similar(u0, T)
        @inbounds @simd ivdep for i in eachindex(u0)
            tmp[i] = u0[i] / sk[i]
        end
    else
        tmp = @.. broadcast = false u0 / sk
    end

    d₀ = internalnorm(tmp, t)

    #=
    Try/catch around the linear solving. This will catch singular matrices defined
    by DAEs and thus we use the _tType(1//10^(6)) default from Hairer. Note that
    this will not always catch singular matrices, an example from Andreas:

    julia> A = fill(rand(), 2, 2)
    2×2 Array{Float64,2}:
     0.637947  0.637947
     0.637947  0.637947

    julia> inv(A)
    2×2 Array{Float64,2}:
      9.0072e15  -9.0072e15
     -9.0072e15   9.0072e15

    The only way to make this more correct is to check

    issingular(A) = rank(A) < min(size(A)...)

    but that would introduce another svdfact in rank (which may not be possible
    anyways if the mass_matrix is not actually an array). Instead we stick to the
    user-chosen factorization. Sometimes this will cause `ftmp` to be absurdly
    large like shown there, but that later gets caught in the quick estimates
    below which then makes it spit out the default

    dt₀ = _tType(1//10^(6))

    so that is a a cheaper way to get to the same place than an svdfact and
    still works for matrix-free definitions of the mass matrix.
    =#

    ftmp = nothing
    if prob.f.mass_matrix != I && (
            !(prob.f isa DynamicalODEFunction) ||
                any(mm != I for mm in prob.f.mass_matrix)
        )
        ftmp = zero(f₀)
        try
            integrator.alg.linsolve(ftmp, copy(prob.f.mass_matrix), f₀, true)
            copyto!(f₀, ftmp)
        catch
            result_dt = tdir * max(smalldt, dtmin)
            @SciMLMessage(
                lazy"Mass matrix appears singular, using default small timestep: dt = $(result_dt)",
                integrator.opts.verbose, :near_singular
            )
            return result_dt
        end
    end

    # d₁: fold in diffusion terms when g !== nothing
    # Pre-initialize g₀ so JET doesn't flag it as potentially undefined in the d₂ block
    g₀ = nothing
    if g !== nothing
        if noise_prototype !== nothing
            g₀ = zero(noise_prototype)
        else
            g₀ = zero(u0)
        end
        g(g₀, u0, p, t)
        g₀ .*= 3
        d₁ = internalnorm(
            max.(internalnorm.(f₀ .+ g₀, t), internalnorm.(f₀ .- g₀, t)) ./ sk, t
        )
    else
        if u0 isa Array
            @inbounds @simd ivdep for i in eachindex(u0)
                tmp[i] = f₀[i] / sk[i] * oneunit_tType
            end
        else
            @.. broadcast = false tmp = f₀ / sk * oneunit_tType
        end
        d₁ = internalnorm(tmp, t)
    end

    # Better than checking any(x->any(isnan, x), f₀)
    # because it also checks if partials are NaN
    # https://discourse.julialang.org/t/incorporating-forcing-functions-in-the-ode-model/70133/26
    if isnan(d₁)
        @SciMLMessage(
            "First function call produced NaNs. Exiting. Double check that none of the initial conditions, parameters, or timespan values are NaN.",
            integrator.opts.verbose, :init_NaN
        )
        return tdir * dtmin
    end

    dt₀ = ifelse(
        (d₀ < 1 // 10^(5)) |
            (d₁ < 1 // 10^(5)), smalldt,
        convert(
            _tType,
            oneunit_tType * SciMLBase.value(
                (d₀ / d₁) /
                    100
            )
        )
    )
    # if d₀ < 1//10^(5) || d₁ < 1//10^(5)
    #   dt₀ = smalldt
    # else
    #   dt₀ = convert(_tType,oneunit_tType*(d₀/d₁)/100)
    # end
    dt₀ = min(dt₀, dtmax_tdir)

    if typeof(one(_tType)) <: AbstractFloat && dt₀ < 10eps(_tType) * oneunit(_tType)
        # This catches Andreas' non-singular example
        # should act like it's singular
        result_dt = tdir * max(smalldt, dtmin)
        @SciMLMessage(
            lazy"Initial timestep too small (near machine epsilon), using default: dt = $(result_dt)",
            integrator.opts.verbose, :dt_epsilon
        )
        return result_dt
    end

    dt₀_tdir = tdir * dt₀

    u₁ = zero(u0) # required by DEDataArray

    if u0 isa Array
        @inbounds @simd ivdep for i in eachindex(u0)
            u₁[i] = u0[i] + dt₀_tdir * f₀[i]
        end
    else
        @.. broadcast = false u₁ = u0 + dt₀_tdir * f₀
    end
    f₁ = zero(f₀)
    f(f₁, u₁, p, t + dt₀_tdir)

    if prob.f.mass_matrix != I && (
            !(prob.f isa DynamicalODEFunction) ||
                any(mm != I for mm in prob.f.mass_matrix)
        )
        integrator.alg.linsolve(ftmp, prob.f.mass_matrix, f₁, false)
        copyto!(f₁, ftmp)
    end

    # Constant zone before callback
    # Just return first guess
    # Avoids AD issues.
    # `==` is not guaranteed to return `Bool` (e.g. PyCall `PyObject` arrays
    # return `Vector{Bool}` — JuliaPy/PyCall.jl#900 / OrdinaryDiffEq.jl#1402).
    # Coerce array-valued equality so the boolean context always receives a Bool.
    length(u0) > 0 && _bool_equal(f₀, f₁) && return tdir * max(dtmin, 100dt₀)

    # d₂: fold in diffusion terms when g !== nothing
    if g !== nothing
        if noise_prototype !== nothing
            g₁ = zero(noise_prototype)
        else
            g₁ = zero(u0)
        end
        g(g₁, u₁, p, t + dt₀_tdir)
        g₁ .*= 3
        ΔgMax = max.(internalnorm.(g₀ .- g₁, t), internalnorm.(g₀ .+ g₁, t))
        d₂ = internalnorm(
            max.(internalnorm.(f₁ .- f₀ .+ ΔgMax, t), internalnorm.(f₁ .- f₀ .- ΔgMax, t)) ./ sk,
            t
        ) / dt₀
    else
        if u0 isa Array
            @inbounds @simd ivdep for i in eachindex(u0)
                tmp[i] = (f₁[i] - f₀[i]) / sk[i] * oneunit_tType
            end
        else
            @.. broadcast = false tmp = (f₁ - f₀) / sk * oneunit_tType
        end
        d₂ = internalnorm(tmp, t) / dt₀ * oneunit_tType
    end
    # Hairer has d₂ = sqrt(sum(abs2,tmp))/dt₀, note the lack of norm correction

    max_d₁d₂ = max(d₁, d₂)
    if max_d₁d₂ <= 1 // Int64(10)^(15)
        dt₁ = max(convert(_tType, oneunit_tType * 1 // 10^(6)), dt₀ * 1 // 10^(3))
    else
        dt₁ = convert(
            _tType,
            oneunit_tType *
                SciMLBase.value(
                10.0^(-(2 + log10(max_d₁d₂)) / order)
            )
        )
    end
    return tdir * max(dtmin, min(100dt₀, dt₁, dtmax_tdir))
end

# ODE iip entry point
"""
    ode_determine_initdt(u0, t, tdir, dtmax, abstol, reltol, internalnorm, prob, integrator) -> dt

Compute an automatic initial step size for `prob` using the standard
Hairer–Nørsett–Wanner heuristic (based on the norms of `f` and its derivative at
`u0`), respecting `dtmax`, the tolerances, and the integration direction `tdir`.
"""
function ode_determine_initdt(
        u0, t, tdir, dtmax, abstol, reltol, internalnorm,
        prob::SciMLBase.AbstractODEProblem{uType, tType, true},
        integrator
    ) where {tType, uType}
    return _ode_initdt_iip(
        u0, t, tdir, dtmax, abstol, reltol, internalnorm,
        prob, nothing, nothing,
        get_current_alg_order(integrator.alg, integrator.cache), integrator
    )
end

const TYPE_NOT_CONSTANT_MESSAGE = """
Detected non-constant types in an out-of-place ODE solve, i.e. for
`du = f(u,p,t)` we see `typeof(du) !== typeof(u/t)`. This is not
supported by OrdinaryDiffEq.jl's solvers. Please either make `f`
type-constant (i.e. typeof(du) === typeof(u/t)) or use the mutating
in-place form `f(du,u,p,t)` (which is type-constant by construction).

Note that one common case for this is when computing with GPUs, using
`Float32` for `u0` and `Float64` for `tspan`. To correct this, ensure
that the element type of `tspan` matches the preferred compute type,
for example `ODEProblem(f,0f0,(0f0,1f0))` for `Float32`-based time.
"""

struct TypeNotConstantError <: Exception
    u0::Type
    f₀::Type
end

function Base.showerror(io::IO, e::TypeNotConstantError)
    println(io, TYPE_NOT_CONSTANT_MESSAGE)
    print(io, "typeof(u/t) = ")
    println(io, e.u0)
    print(io, "typeof(du) = ")
    return println(io, e.f₀)
end

@muladd function _ode_initdt_oop(
        u0, t, tdir, dtmax, abstol, reltol, internalnorm,
        prob, g, order, integrator
    )
    _tType = eltype(t)
    f = prob.f
    p = prob.p
    oneunit_tType = oneunit(t)
    dtmax_tdir = tdir * dtmax

    dtmin = nextfloat(max(integrator.opts.dtmin, convert(_tType, oneunit_tType * eps(SciMLBase.value(t)))))
    smalldt = max(dtmin, convert(_tType, oneunit_tType * 1 // 10^(6)))

    if integrator.isdae
        return tdir * max(smalldt, dtmin)
    end

    # With zero continuous states there is nothing to estimate from, so fall back
    # to a finite default dt rather than indexing into the empty state.
    if isempty(u0)
        return tdir * max(smalldt, dtmin)
    end

    sk = @.. broadcast = false abstol + internalnorm(u0, t) * reltol
    d₀ = internalnorm(u0 ./ sk, t)

    f₀ = f(u0, p, t)

    # Use the overloadable DiffEqBase.NAN_CHECK hook (same intent as the IIP
    # isnan(d₁) path) rather than nested any(isnan, ·), which custom array /
    # field types cannot sensibly overload (OrdinaryDiffEq #1404).
    if DiffEqBase.NAN_CHECK(f₀)
        @SciMLMessage(
            "First function call produced NaNs. Exiting. Double check that none of the initial conditions, parameters, or timespan values are NaN.",
            integrator.opts.verbose, :init_NaN
        )
        return tdir * dtmin
    end

    inferredtype = Base.promote_op(/, typeof(u0), typeof(oneunit(t)))
    if !(f₀ isa inferredtype)
        throw(TypeNotConstantError(inferredtype, typeof(f₀)))
    end

    # d₁: fold in diffusion terms when g !== nothing
    # Pre-initialize g₀ so JET doesn't flag it as potentially undefined in the d₂ block
    g₀ = nothing
    if g !== nothing
        g₀ = 3g(u0, p, t)
        if DiffEqBase.NAN_CHECK(g₀)
            @SciMLMessage(
                "First function call for g produced NaNs. Exiting.",
                integrator.opts.verbose, :init_NaN
            )
            return tdir * dtmin
        end
        d₁ = internalnorm(
            max.(internalnorm.(f₀ .+ g₀, t), internalnorm.(f₀ .- g₀, t)) ./ sk, t
        )
    else
        d₁ = internalnorm(f₀ ./ sk .* oneunit_tType, t)
    end

    # Also catch NaN AD partials that NAN_CHECK on values may miss (matches IIP).
    if isnan(d₁)
        @SciMLMessage(
            "First function call produced NaNs. Exiting. Double check that none of the initial conditions, parameters, or timespan values are NaN.",
            integrator.opts.verbose, :init_NaN
        )
        return tdir * dtmin
    end

    if d₀ < 1 // 10^(5) || d₁ < 1 // 10^(5)
        dt₀ = smalldt
    else
        dt₀ = convert(_tType, oneunit_tType * SciMLBase.value((d₀ / d₁) / 100))
    end
    dt₀ = min(dt₀, dtmax_tdir)
    dt₀_tdir = tdir * dt₀

    u₁ = @.. broadcast = false u0 + dt₀_tdir * f₀
    f₁ = f(u₁, p, t + dt₀_tdir)

    # Constant zone before callback
    # Just return first guess
    # Avoids AD issues.
    # See iip path: coerce array-valued `==` (OrdinaryDiffEq.jl#1402).
    _bool_equal(f₀, f₁) && return tdir * max(dtmin, 100dt₀)

    # d₂: fold in diffusion terms when g !== nothing
    if g !== nothing
        g₁ = 3g(u₁, p, t + dt₀_tdir)
        ΔgMax = max.(internalnorm.(g₀ .- g₁, t), internalnorm.(g₀ .+ g₁, t))
        d₂ = internalnorm(
            max.(internalnorm.(f₁ .- f₀ .+ ΔgMax, t), internalnorm.(f₁ .- f₀ .- ΔgMax, t)) ./ sk,
            t
        ) / dt₀
    else
        d₂ = internalnorm((f₁ .- f₀) ./ sk .* oneunit_tType, t) / dt₀ * oneunit_tType
    end

    max_d₁d₂ = max(d₁, d₂)
    if max_d₁d₂ <= 1 // Int64(10)^(15)
        dt₁ = max(smalldt, dt₀ * 1 // 10^(3))
    else
        dt₁ = _tType(
            oneunit_tType *
                SciMLBase.value(
                10^(-(2 + log10(max_d₁d₂)) / order)
            )
        )
    end
    return tdir * max(dtmin, min(100dt₀, dt₁, dtmax_tdir))
end

# ODE oop entry point
function ode_determine_initdt(
        u0, t, tdir, dtmax, abstol, reltol, internalnorm,
        prob::SciMLBase.AbstractODEProblem{uType, tType, false},
        integrator
    ) where {uType, tType}
    return _ode_initdt_oop(
        u0, t, tdir, dtmax, abstol, reltol, internalnorm,
        prob, nothing,
        get_current_alg_order(integrator.alg, integrator.cache), integrator
    )
end

@inline function ode_determine_initdt(
        u0, t, tdir, dtmax, abstol, reltol, internalnorm,
        prob::SciMLBase.AbstractDAEProblem{
            duType, uType,
            tType,
        },
        integrator
    ) where {duType, uType, tType}
    _tType = eltype(tType)
    tspan = prob.tspan
    init_dt = abs(tspan[2] - tspan[1])
    init_dt = isfinite(init_dt) ? init_dt : oneunit(_tType)
    # Must include tdir so reversed tspans get a negative initial dt
    # (DifferentialEquations.jl #908).
    return convert(_tType, tdir * init_dt * 1 // 10^(6))
end

# RODE/SDE iip entry: folds noise terms into the Hairer-Wanner estimate
function ode_determine_initdt(
        u0, t, tdir, dtmax, abstol, reltol, internalnorm,
        prob::SciMLBase.AbstractRODEProblem{uType, tType, true},
        order, integrator
    ) where {tType, uType}
    if _get_P(integrator) !== nothing
        return tdir * dtmax / 1.0e6
    end
    g = prob.f.g
    noise_proto = hasproperty(prob, :noise_rate_prototype) ? prob.noise_rate_prototype :
        nothing
    effective_order = g !== nothing ? order + 1 // 2 : order
    return _ode_initdt_iip(
        u0, t, tdir, dtmax, abstol, reltol, internalnorm,
        prob, g, noise_proto, effective_order, integrator
    )
end

# RODE/SDE oop entry: folds noise terms into the Hairer-Wanner estimate
function ode_determine_initdt(
        u0, t, tdir, dtmax, abstol, reltol, internalnorm,
        prob::SciMLBase.AbstractRODEProblem{uType, tType, false},
        order, integrator
    ) where {tType, uType}
    if _get_P(integrator) !== nothing
        return tdir * dtmax / 1.0e6
    end
    g = prob.f.g
    effective_order = g !== nothing ? order + 1 // 2 : order
    return _ode_initdt_oop(
        u0, t, tdir, dtmax, abstol, reltol, internalnorm,
        prob, g, effective_order, integrator
    )
end
