using Base: @propagate_inbounds

@doc raw"""
    integrate(integrand::Function, f::AbstractFilament, i::Int, quad::AbstractQuadrature; limits = nothing)
    integrate(integrand::Function, s::Segment, quad::AbstractQuadrature; limits = nothing)

Estimate integral along a filament segment using the chosen quadrature.

Integration is performed in the segment `(f[i], f[i + 1])`.

The function `integrand(Î¶)` takes the arguments `f`, `i` and ``Î¶ âˆˆ [0, 1]``.
The first two are a bit redundant since they're the same filament and node index passed to
this function, while `Î¶` corresponds to the position of a point within the segment.
See further below for some examples.

By default, the integration is performed along the whole segment, that is, for the range ``Î¶ âˆˆ [0, 1]``.
It is possible to integrate over a subset of the segment, i.e. for ``Î¶ âˆˆ [a, b]`` with ``0 â‰¤ a â‰¤ b â‰¤ 1``.
For this, one should pass the keyword argument `limits = (a, b)`.

# Examples

Estimate arc length of segment ``[i, i + 1]``, given by ``â„“ = âˆ«_{t_{i}}^{t_{i + 1}} |âˆ‚_t \bm{X}(t)| \, \mathrm{d}t``:

```julia
quad = GaussLegendre(4)  # quadrature rule
â„“ = integrate(f, i, quad) do f, i, Î¶
    norm(f(i, Î¶, Derivative(1)))  # = |âˆ‚â‚œğ��—|
end
```

Alternatively:
```julia
quad = GaussLegendre(4)  # quadrature rule
s = Segment(f, i)
â„“ = integrate(s, quad) do s, Î¶
    norm(s(Î¶, Derivative(1)))  # = |âˆ‚â‚œğ��—|
end
```
"""
function integrate(
        integrand::F, f::AbstractFilament, i::Int, quad::AbstractQuadrature;
        limits = nothing,
        _args = (f, i),  # integrand arguments; used internally
    ) where {F}
    _integrate(integrand, limits, f, i, quad; _args)
end

# Here `limits` can be either:
# - `nothing`, for default limits [0, 1]
# - a tuple (a, b), with 0 â‰¤ a < b â‰¤ 1.
function _integrate(integrand::F, limits, f, i, quad; _args) where {F}
    ts = knots(f)
    T = eltype(ts)
    Î”t = ts[i + 1] - ts[i]
    # We integrate wrt Î¶ âˆˆ [0, 1], then rescale by Î”t to obtain an integral wrt t.
    lims = to_lims_arg(T, limits)
    Î”t * Quadratures.integrate(quad, lims) do Î¶
        @inline
        @inline integrand(_args..., Î¶)
    end
end

# Convert limits to `lims` arg of Quadratures.integrate.
to_lims_arg(::Type{T}, ::Nothing) where {T} = T  # default limits (0, 1)
to_lims_arg(::Type{T}, (a, b)::NTuple{2, Real}) where {T} = (T(a), T(b))

function integrate(integrand::F, s::Segment, quad::AbstractQuadrature; kws...) where {F}
    integrate(integrand, s.f, s.i, quad; _args = (s,), kws...)
end

"""
    integrate(integrand::Function, f::AbstractFilament, quad::AbstractQuadrature; limits = nothing)

Estimate integral over a whole filament.

The integral is computed as the sum of the integrals over each filament segment.
The given quadrature rule is applied over each segment, meaning that the total number of
evaluations is given by the number of segments multiplied by the length of the quadrature
rule.

The signature of the integrated function must be
`integrand(f::AbstractFilament, i::Int, Î¶::Real)`, where `i` is the index of the segment of
interest. See below for some examples.

## Examples

Estimate the total length of a closed filament, ``L = âˆ® |âˆ‚_t \\bm{X}(t)| \\, \\mathrm{d}t``:

```julia
quad = GaussLegendre(4)  # quadrature rule

integrand(f, i, Î¶) = norm(f(i, Î¶, Derivative(1)))  # = |âˆ‚â‚œğ��—|

# Here `f` is an existent filament.
L = integrate(integrand, f, quad)

# Or, more conveniently, using a `do` block to define an anonymous function.
L = integrate(f, quad) do ff, i, Î¶
    norm(ff(i, Î¶, Derivative(1)))  # = |âˆ‚â‚œğ��—|
end
```
"""
function integrate(integrand::F, f::AbstractFilament, quad::AbstractQuadrature) where {F}
    sum(segments(f)) do seg
        integrate(integrand, seg.f, seg.i, quad)
    end
end

## ====================================================================================== ##
#  No quadrature case: the idea is to only evaluate quantities on filament nodes,
#  avoiding any interpolation.
## ====================================================================================== ##

# This is an internal type which is just used to replace the interpolation syntax
# f(i, Î¶, args...) by the on-node evaluation syntax f[i, args...].
# When integrating, we consider that we evaluate in the middle of each segment (at Î¶ = 1/2).
# Moreover, the location and the derivative at that point are approximated from the
# locations at the two extremities of the segment, according to a straight-segment
# approximation.
struct NoQuadFilament{
        T,
        F <: AbstractFilament{T},
    } <: AbstractFilament{T, Nothing, AbstractParametrisation}
    f :: F
end

@propagate_inbounds function (u::NoQuadFilament)(i::Int, Î¶::Number)
    u(i, Î¶, Derivative(0))
end

@propagate_inbounds function (u::NoQuadFilament)(i::Int, Î¶::Number, ::Derivative{0})
    (; f,) = u
    # Note: precision could be improved by shifting the "straight" segment in the direction
    # opposite to the curvature vector.
    # See McGreivy et al, Phys. Plasmas 28, 082111 (2021) [https://doi.org/10.1063/5.0058014].
    # However estimating curvatures increases the computation cost, and we want to keep
    # NoQuadrature as cheap as possible.
    (1 - Î¶) * f[i] + Î¶ * f[i + 1]
end

@propagate_inbounds function (u::NoQuadFilament)(i::Int, Î¶::Number, ::Derivative{1})
    (; f,) = u
    ts = knots(f)
    (f[i + 1] - f[i]) ./ (ts[i + 1] - ts[i])
end

# Not sure if this is used...
@propagate_inbounds function (u::NoQuadFilament)(i, Î¶::Number, d)
    (1 - Î¶) * u.f[i, d] + Î¶ * u.f[i + 1, d]
end

function integrate(
        integrand::F, f::AbstractFilament, i::Int, quad::NoQuadrature;
        limits = nothing,
        _args = (NoQuadFilament(f), i),
    ) where {F}
    ts = knots(f)
    args = _noquad_replace_args(_args...)  # replaces filaments by NoQuadFilament
    Î”t = ts[i + 1] - ts[i]
    if limits !== nothing
        a, b = limits
        Î”t *= b - a
        Î¶ = (a + b) / 2
    else
        Î¶ = 0.5
    end
    Î”t * integrand(args..., Î¶)
end

# These are the only 2 possible options for _args:
_noquad_replace_args(f::NoQuadFilament, i::Int) = (f, i)
_noquad_replace_args(s::Segment) = (Segment(NoQuadFilament(s.f), s.i),)  # replace AbstractFilament by NoQuadFilament
