# Element-wise scalar transforms between bounded and unbounded spaces.
# Each function handles 4 cases based on finiteness of lb/ub.

function _to_unbounded(u, lb, ub)
    has_lb = isfinite(lb)
    has_ub = isfinite(ub)
    if has_lb && has_ub
        return logit((u - lb) / (ub - lb))
    elseif has_lb
        return log(u - lb)
    elseif has_ub
        return log(ub - u)
    else
        return u
    end
end

function _from_unbounded(t, lb, ub)
    has_lb = isfinite(lb)
    has_ub = isfinite(ub)
    if has_lb && has_ub
        return lb + (ub - lb) * logistic(t)
    elseif has_lb
        return lb + exp(t)
    elseif has_ub
        return ub - exp(t)
    else
        return t
    end
end

# Clamp a value into the strict interior of [lb, ub] so that _to_unbounded (logit)
# doesn't receive 0 or 1, which would give ±Inf. Only applied once to u0 before
# the initial transform — it's a no-op if u0 is already in the interval.
#
# We use eps^(3/4) (~1.8e-12 for Float64) as the relative nudge factor. Plain eps
# (~2.2e-16) is so small that nudged values can round back to the boundary, while
# eps^(3/4) gives comfortable room without meaningfully changing the starting point.
function _clamp_to_bounds(u, lb, ub)
    has_lb = isfinite(lb)
    has_ub = isfinite(ub)
    eps_frac = eps(typeof(u))^(3 / 4)
    if has_lb && has_ub
        # Margin scales with interval width
        margin = (ub - lb) * eps_frac
        return clamp(u, lb + margin, ub - margin)
    elseif has_lb
        # max(abs(lb), 1) provides a scale factor so the nudge isn't zero when lb == 0
        return max(u, lb + eps_frac * max(abs(lb), one(lb)))
    elseif has_ub
        return min(u, ub - eps_frac * max(abs(ub), one(ub)))
    else
        return u
    end
end

# Normalize bounds: convert nothing to ±Inf vectors, broadcast scalars to match u0.
function _normalize_bound(bound, fill_value, u0)
    T = eltype(u0)

    return if isnothing(bound)
        fill(T(fill_value), size(u0))
    elseif bound isa Number
        fill(T(bound), size(u0))
    else
        T.(bound)
    end
end

function _normalize_bounds(lb, ub, u0)
    new_lb = _normalize_bound(lb, -Inf, u0)
    new_ub = _normalize_bound(ub, Inf, u0)
    return new_lb, new_ub
end

# Wrapper that contains the bounds and a cache to use for storing the
# transformed bounds.
@concrete struct BoundedWrapper{isinplace}
    f
    lb
    ub
    u_cache
end

function _transform_u(w::BoundedWrapper, u)
    tmp = w.u_cache isa FixedSizeDiffCache ? get_tmp(w.u_cache, u) : w.u_cache
    tmp .= _from_unbounded.(u, w.lb, w.ub)
    return tmp
end

function (w::BoundedWrapper{false})(u, p)
    transformed_u = _transform_u(w, u)
    return w.f(transformed_u, p)
end

function (w::BoundedWrapper{true})(resid, u, p)
    transformed_u = _transform_u(w, u)
    w.f(resid, transformed_u, p)
    return resid
end

SciMLBase.isinplace(w::BoundedWrapper{iip}) where {iip} = iip

# Check if bounds transform is needed for a given problem and algorithm.
function needs_bounds_transform(_prob, alg)
    return (
        _prob isa SciMLBase.NonlinearProblem ||
            _prob isa SciMLBase.NonlinearLeastSquaresProblem
    ) &&
        (hasfield(typeof(_prob), :lb) && hasfield(typeof(_prob), :ub)) &&
        (_prob.lb !== nothing || _prob.ub !== nothing) &&
        (isnothing(alg) || !SciMLBase.allowsbounds(alg))
end

# Wrap a problem function with bounds into a BoundedWrapper with no bounds. In a
# nutshell, we transform a parameter `p` with bounds `lb` and `ub` into an
# unbounded parameter `t` using the logistic function to map all values of `t`
# into the interval (lb, ub).
function transform_bounded_problem(prob, alg)
    # `initialization_data` (e.g. from ModelingToolkit) is defined in the original,
    # bounded coordinates. It must run *before* the transform so the resulting `u0`/`p`
    # are consistent; running it on the transformed problem would interpret the unbounded
    # iterate `t` as a bounded state and corrupt the solution. We run it here and strip it
    # from the transformed function below so it isn't re-run in `t`-space.
    prob = run_bounded_initialization(prob, alg)

    lb, ub = _normalize_bounds(prob.lb, prob.ub, prob.u0)

    # Clamp u0 into the interior of the bounds so that _to_unbounded doesn't hit log(0)
    # or log(negative). We nudge by a small fraction of the interval width.
    u0_clamped = _clamp_to_bounds.(prob.u0, lb, ub)
    u0_transformed = _to_unbounded.(u0_clamped, lb, ub)

    # PreallocationTools is only supported by ForwardDiff so we only use
    # FixedSizeDiffCache if we're using ForwardDiff. Not every algorithm has an
    # `autodiff` field (e.g. `QuasiNewtonAlgorithm`), so guard the access.
    alg_ad = alg !== nothing && hasproperty(alg, :autodiff) ? alg.autodiff : nothing
    u_cache = if alg_ad === nothing || alg_ad isa AutoForwardDiff
        FixedSizeDiffCache(prob.u0)
    else
        similar(prob.u0)
    end

    orig_f = prob.f
    # Unwrap AutoSpecializeCallable before wrapping in BoundedWrapper.
    # BoundedWrapper transforms arguments, so the FunctionWrapper signatures won't match.
    unwrapped_orig_f = if is_fw_wrapped(orig_f.f)
        @set orig_f.f = get_raw_f(orig_f.f)
    else
        orig_f
    end
    wrapped = BoundedWrapper{SciMLBase.isinplace(prob)}(unwrapped_orig_f, lb, ub, u_cache)

    new_f = if orig_f isa NonlinearFunction
        @set orig_f.f = wrapped
    else
        wrapped
    end

    transformed_prob = remake(
        prob; f = new_f, u0 = u0_transformed, lb = nothing, ub = nothing,
        build_initializeprob = Val{false}
    )

    return transformed_prob
end

# Run problem initialization (if any) in the original bounded coordinates so that the
# resulting `u0`/`p` are consistent before the bounds transform is applied. Returns the
# problem unchanged when there is no `OverrideInitData` to run.
function run_bounded_initialization(prob, alg)
    SciMLBase.has_initialization_data(prob.f) || return prob
    prob.f.initialization_data isa SciMLBase.OverrideInitData || return prob
    iip = SciMLBase.isinplace(prob)
    u0, p, _ = SciMLBase.get_initial_values(
        prob, prob, prob.f, SciMLBase.OverrideInit(), Val(iip);
        nlsolve_alg = alg, abstol = get_abstol(prob), reltol = get_reltol(prob)
    )
    return remake(prob; u0, p, build_initializeprob = Val{false})
end
