"""
    ExtrapolationController(; kwargs...) -> ExtrapolationController
    ExtrapolationController(::Type{QT}, alg; kwargs...) -> ExtrapolationController

Internal step-size controller for the extrapolation algorithms. OrdinaryDiffEq selects
this controller automatically; solver users should not construct or dispatch on it.

!!! warning "Developer implementation"
    This type is implementation machinery, not user-facing API. Its representation and
    direct construction may change as the extrapolation controller evolves.

# Arguments

- `QT::Type`: scalar type used to resolve controller options.
- `alg`: extrapolation algorithm whose default controller options are used.

# Keywords

- `qmin`: lower bound on the step-size reduction factor.
- `qmax`: upper bound on the step-size growth factor.
- `qmax_first_step`: upper growth bound after the first accepted step.
- `gamma`: safety factor applied to a proposed step-size change.
- `qsteady_min`: lower edge of the interval in which the step size is unchanged.
- `qsteady_max`: upper edge of the interval in which the step size is unchanged.
- `failfactor`: step-size reduction factor after a nonlinear-solver failure.
- `discontinuity_detection`: enable autonomous discontinuity detection when `true`.

Omitted keywords use the defaults associated with `alg` when the controller cache is
created.

# Returns

- `ExtrapolationController`: unresolved controller when `QT` and `alg` are omitted, or
  a controller whose common options have been resolved to `QT` otherwise.

# Examples

```julia
import OrdinaryDiffEqExtrapolation
using OrdinaryDiffEqExtrapolation: ExtrapolationMidpointDeuflhard

alg = ExtrapolationMidpointDeuflhard()
controller = OrdinaryDiffEqExtrapolation.ExtrapolationController(Float64, alg)
```
"""
struct ExtrapolationController{B <: Union{NamedTuple, OrdinaryDiffEqCore.CommonControllerOptions}} <: AbstractController
    basic::B
end

ExtrapolationController(; kwargs...) = ExtrapolationController(NamedTuple(kwargs))

function ExtrapolationController(::Type{QT}, alg; kwargs...) where {QT}
    return ExtrapolationController(
        OrdinaryDiffEqCore.resolve_basic(NamedTuple(kwargs), alg, QT),
    )
end

mutable struct ExtrapolationControllerCache{QT, E, NLPType} <: AbstractControllerCache
    controller::ExtrapolationController{OrdinaryDiffEqCore.CommonControllerOptions{QT, NLPType}}
    beta1::QT
    gamma::QT
    qold::QT
    EEst::E
end

function setup_controller_cache(alg, cache, controller::ExtrapolationController, ::Type{E}, disco_probs) where {E}
    QT = _controller_scalar_type(controller.basic)
    basic = OrdinaryDiffEqCore.resolve_basic(controller.basic, alg, QT; disco_probs)
    resolved = ExtrapolationController(basic)
    T = QT
    return ExtrapolationControllerCache{T, E, eltype(disco_probs)}(
        resolved, T(1), T(1), T(1 // 10^4), oneunit(E),
    )
end

function reset_alg_dependent_opts!(controller::ExtrapolationControllerCache, alg1, alg2)
    if controller.beta1 == beta1_default(alg1, beta2_default(alg1))
        controller.beta1 = beta1_default(alg2, beta2_default(alg2))
    end
    return nothing
end

function OrdinaryDiffEqCore.sync_controllers!(cache1::ExtrapolationControllerCache, cache2::ExtrapolationControllerCache)
    cache1.beta1 = cache2.beta1
    cache1.gamma = cache2.gamma
    cache1.qold = cache2.qold
    return nothing
end

stepsize_controller_internal!(integrator, alg) = stepsize_controller_internal!(integrator, integrator.controller_cache, alg)
stepsize_predictor!(integrator, alg, n_new) = stepsize_predictor!(integrator, integrator.controller_cache, alg, n_new)

@inline function stepsize_controller!(
        integrator,
        cache::ExtrapolationControllerCache,
        alg::Union{
            ExtrapolationMidpointDeuflhard,
            ExtrapolationMidpointHairerWanner,
            ImplicitDeuflhardExtrapolation,
            ImplicitHairerWannerExtrapolation,
            ImplicitEulerExtrapolation,
            ImplicitEulerBarycentricExtrapolation,
        }
    )
    # Dummy function
    # ExtrapolationMidpointDeuflhard's stepsize scaling is stored in the cache;
    # it is computed by  stepsize_controller_internal! (in perform_step!) resp. stepsize_predictor!
    # (in step_accept_controller! and step_reject_controller!)
    return zero(typeof(get_qmax(integrator)))
end

function stepsize_controller_internal!(
        integrator,
        cache::ExtrapolationControllerCache,
        alg::Union{
            ExtrapolationMidpointDeuflhard,
            ImplicitDeuflhardExtrapolation,
        }
    )
    # Standard step size controller
    # Compute and save the stepsize scaling based on the latest error estimate of the current order
    (; controller) = cache

    if iszero(OrdinaryDiffEqCore.get_EEst(integrator))
        q = inv(get_qmax(integrator))
    else
        # Update gamma and beta1
        cache.beta1 = typeof(cache.beta1)(1 // (2integrator.cache.n_curr + 1))
        cache.gamma = fastpower(
            typeof(cache.gamma)(1 // 4),
            cache.beta1
        )
        # Compute new stepsize scaling
        qtmp = fastpower(OrdinaryDiffEqCore.get_EEst(integrator), cache.beta1) /
            cache.gamma
        @fastmath q = max(inv(get_qmax(integrator)), min(inv(get_qmin(integrator)), qtmp))
    end
    return integrator.cache.Q[integrator.cache.n_curr - alg.min_order + 1] = q
end

function stepsize_predictor!(
        integrator,
        cache::ExtrapolationControllerCache,
        alg::Union{
            ExtrapolationMidpointDeuflhard,
            ImplicitDeuflhardExtrapolation,
        }, n_new::Int
    )
    # Compute and save the stepsize scaling for order n_new based on the latest error estimate of the current order.
    (; controller) = cache

    if iszero(OrdinaryDiffEqCore.get_EEst(integrator))
        q = inv(get_qmax(integrator))
    else
        # Initialize
        t = integrator.t
        EEst = OrdinaryDiffEqCore.get_EEst(integrator)
        (; stage_number) = integrator.cache
        tol = integrator.opts.internalnorm(integrator.opts.reltol, t) # Deuflhard's approach relies on EEstD ≈ ||relTol||
        s_curr = stage_number[integrator.cache.n_curr - alg.min_order + 1]
        s_new = stage_number[n_new - alg.min_order + 1]
        # Update gamma and beta1
        cache.beta1 = typeof(cache.beta1)(1 // (2integrator.cache.n_curr + 1))
        cache.gamma = fastpower(
            typeof(cache.gamma)(1 // 4),
            cache.beta1
        )
        # Compute new stepsize scaling
        qtmp = EEst *
            fastpower(
            fastpower(tol, (1.0 - s_curr / s_new)),
            cache.beta1
        ) / cache.gamma
        @fastmath q = max(inv(get_qmax(integrator)), min(inv(get_qmin(integrator)), qtmp))
    end
    return integrator.cache.Q[n_new - alg.min_order + 1] = q
end

function step_accept_controller!(
        integrator,
        alg::Union{
            ExtrapolationMidpointDeuflhard,
            ImplicitDeuflhardExtrapolation,
        }, q
    )
    # Compute new order and stepsize, return new stepsize
    (; min_order, max_order) = alg
    (; n_curr, n_old, Q) = integrator.cache
    s = integrator.cache.stage_number

    # Compute new order based on available quantities
    tmp = (min_order:n_curr) .- min_order .+ 1 # Index range of quantities computed so far
    dt_new = Vector{eltype(Q)}(undef, length(tmp) + 1)
    dt_new[1:(end - 1)] = integrator.dt ./ Q[tmp] # Store for the possible new stepsizes
    dtmin = timedepentdtmin(integrator)
    dt_new[1:(end - 1)] = max.(
        dtmin,
        min.(abs(integrator.opts.dtmax), abs.(dt_new[1:(end - 1)]))
    ) # Safety scaling

    # n_new is the most efficient order of the last step
    work = s[tmp] ./ dt_new[1:(end - 1)]
    n_new = argmin(work) + min_order - 1

    # Check if n_new may be increased
    if n_new == n_curr < min(max_order, n_old + 1) # cf. win_max in perform_step! of the last step
        # Predict stepsize scaling for order (n_new + 1)
        stepsize_predictor!(integrator, alg, n_new + 1) # Update cache.Q

        # Compute and scale the corresponding stepsize
        dt_new[end] = integrator.dt ./ Q[tmp[end] + 1]
        dt_new[end] = max(dtmin, min(abs(integrator.opts.dtmax), abs.(dt_new[end])))

        # Check if (n_new  + 1) would have been more efficient than n_new
        if work[end] > s[tmp[end] + 1] / dt_new[end]
            n_new = n_new + 1
        end
    end

    integrator.cache.n_curr = n_new
    return dt_new[n_new - min_order + 1]
end

function step_reject_controller!(
        integrator,
        alg::Union{
            ExtrapolationMidpointDeuflhard,
            ImplicitDeuflhardExtrapolation,
        }
    )
    # Compute and save reduced stepsize dt_red of order n_old
    # Use the latest error estimate to predict dt_red if an estimate of order n_old is not available
    if integrator.cache.n_curr < integrator.cache.n_old
        stepsize_predictor!(integrator, alg, integrator.cache.n_old) # Update cache.Q
    end
    integrator.cache.n_curr = integrator.cache.n_old # Reset order for redoing the rejected step
    dt_red = integrator.dt /
        integrator.cache.Q[integrator.cache.n_old - integrator.alg.min_order + 1]
    dtmin = timedepentdtmin(integrator)
    dt_red = integrator.tdir * max(dtmin, min(abs(integrator.opts.dtmax), abs(dt_red))) # Safety scaling
    return integrator.dt = dt_red
end

function stepsize_controller_internal!(
        integrator,
        cache::ExtrapolationControllerCache,
        alg::Union{
            ExtrapolationMidpointHairerWanner,
            ImplicitHairerWannerExtrapolation,
            ImplicitEulerExtrapolation,
            ImplicitEulerBarycentricExtrapolation,
        }
    )
    # Standard step size controller
    # Compute and save the stepsize scaling based on the latest error estimate of the current order
    (; controller) = cache

    return if alg isa
            Union{
            ImplicitEulerExtrapolation, ImplicitEulerBarycentricExtrapolation,
            ImplicitHairerWannerExtrapolation,
        }
        if iszero(OrdinaryDiffEqCore.get_EEst(integrator))
            q = inv(get_qmax(integrator))
        else
            # Update gamma and beta1
            if alg isa ImplicitHairerWannerExtrapolation
                cache.beta1 = typeof(cache.beta1)(
                    1 //
                        (2integrator.cache.n_curr + 1)
                )
            elseif alg isa ImplicitEulerExtrapolation
                cache.beta1 = typeof(cache.beta1)(1 // (integrator.cache.n_curr))
            else
                cache.beta1 = typeof(cache.beta1)(
                    1 //
                        (integrator.cache.n_curr - 1)
                )
            end
            cache.gamma = fastpower(
                typeof(cache.gamma)(
                    65 //
                        100
                ),
                cache.beta1
            )
            # Compute new stepsize scaling
            qtmp = fastpower(OrdinaryDiffEqCore.get_EEst(integrator), cache.beta1) /
                (cache.gamma)
            @fastmath q = max(
                inv(get_qmax(integrator)),
                min(inv(get_qmin(integrator)), qtmp)
            )
        end
        integrator.cache.Q[integrator.cache.n_curr + 1] = q
    else
        if iszero(OrdinaryDiffEqCore.get_EEst(integrator))
            q = inv(get_qmax(integrator))
        else
            # Update gamma and beta1
            cache.beta1 = typeof(cache.beta1)(1 // (2integrator.cache.n_curr + 1))
            cache.gamma = fastpower(
                typeof(cache.gamma)(
                    65 //
                        100
                ),
                cache.beta1
            )
            # Compute new stepsize scaling
            qtmp = fastpower(OrdinaryDiffEqCore.get_EEst(integrator), cache.beta1) /
                cache.gamma
            @fastmath q = max(
                inv(get_qmax(integrator)),
                min(inv(get_qmin(integrator)), qtmp)
            )
        end
        integrator.cache.Q[integrator.cache.n_curr + 1] = q
    end
end

function step_accept_controller!(
        integrator,
        alg::Union{
            ExtrapolationMidpointHairerWanner,
            ImplicitHairerWannerExtrapolation,
            ImplicitEulerExtrapolation,
            ImplicitEulerBarycentricExtrapolation,
        }, q
    )
    # Compute new order and stepsize, return new stepsize
    (; min_order, max_order) = alg
    (; n_curr, n_old, Q, sigma, work, dt_new) = integrator.cache
    s = integrator.cache.stage_number

    # Compute new order based on available quantities
    win_min_old = min(n_old, n_curr) - 1 # cf. win_min in perform_step! of the last step
    tmp = win_min_old:(max(n_curr, n_old) + 1) # Index range for the new order
    fill!(dt_new, zero(eltype(dt_new)))
    @.. broadcast = false Q = integrator.dt / Q
    copyto!(dt_new, win_min_old, Q, win_min_old, (max(n_curr, n_old) + 1) - win_min_old + 1)
    @.. broadcast = false Q = integrator.dt / Q
    dtmin = timedepentdtmin(integrator)
    fill!(work, zero(eltype(work))) # work[n] is the work for order (n-1)
    for i in tmp
        work[i] = s[i] / dt_new[i]
    end
    # Order selection
    n_new = n_old
    if n_curr == min_order # Enforce min_order + 1 ≦ n_new
        n_new = min_order + 1
    else
        if n_curr <= n_old
            if work[n_curr - 1] < sigma * work[n_curr]
                n_new = max(n_curr - 1, n_old - 1, min_order + 1) # Enforce min_order + 1≦ n_new
            elseif work[n_curr] < sigma * work[n_curr - 1]
                n_new = min(n_curr + 1, max_order - 1) # Enforce n_new ≦ max_order - 1
            else
                n_new = n_curr # min_order + 1 ≦ n_curr
            end
        else
            if work[n_old] < sigma * work[n_old + 1]
                n_new = max(n_old - 1, min_order + 1)  # Enforce min_order + 1 ≦ n_new
            end
            if work[n_curr + 1] < sigma * work[n_new + 1]
                n_new = min(n_new + 1, max_order - 1) # Enforce n_new ≦ max_order - 1
            end
        end
    end
    integrator.cache.n_curr = n_new

    # Stepsize selection
    if n_new == n_curr + 1
        # Compute the new stepsize of order n_new based on the optimal stepsize of order n_curr
        dt_new[n_new + 1] = s[n_curr + 2] / s[n_curr + 1] * dt_new[n_curr + 1]
        dt_new[n_new + 1] = max(
            dtmin,
            min(abs(integrator.opts.dtmax), abs(dt_new[n_new + 1]))
        )
    end
    return dt_new[n_new + 1]
end

function step_reject_controller!(
        integrator,
        alg::Union{
            ExtrapolationMidpointHairerWanner,
            ImplicitHairerWannerExtrapolation,
            ImplicitEulerExtrapolation,
            ImplicitEulerBarycentricExtrapolation,
        }
    )
    # Compute and save order and stepsize for redoing the current step
    (; n_old, n_curr, Q) = integrator.cache

    # Order selection
    n_red = n_old
    if n_curr == n_old - 1
        n_red = max(alg.min_order + 1, n_old - 1) # Enforce min_order + 1 ≦ n_red
    end
    integrator.cache.n_curr = n_red

    # Stepsize selection
    dt_red = integrator.dt / Q[n_red + 1]
    dtmin = timedepentdtmin(integrator)
    dt_red = integrator.tdir * max(dtmin, min(abs(integrator.opts.dtmax), abs(dt_red))) # Safety scaling
    return integrator.dt = dt_red
end
