import Base: first, last, length,
    getindex, lastindex, eachindex, iterate,
    eltype, similar,
    show

# * AbstractKnotSet
"""
    AbstractKnotSet{k,ml,mr,T}

Abstract base for B-spline knot sets. `T` is the `eltype` of the knot
set, `k` is the order of the piecewise polynomials (`order = degree +
1`) and `ml` and `mr` are the knot multiplicities of the left and
right endpoints, respectively.  """
abstract type AbstractKnotSet{k,ml,mr,T} <: AbstractVector{T} end

distribution(::AbstractKnotSet) = NonUniform()

"""
    assert_multiplicities(k,ml,mr,t)

Assert that the multiplicities at the endpoints, `ml` and `mr`,
respectively, are consistent with the order `k`. Also check that the
amount of knots in the knot set `t` are enough to support the
requested order `k`.
"""
function assert_multiplicities(k,ml,mr,t)
    k > 0 || throw(ArgumentError("Polynomial order must be a positive integer, got k = $k"))
    ml ∈ 1..k || throw(ArgumentError("Left multiplicity has to be in range 1..$(k), got $ml"))
    mr ∈ 1..k || throw(ArgumentError("Right multiplicity has to be in range 1..$(k), got $mr"))
    n = length(t)
    nreq = max(2, k + 3 - ml - mr)
    n ≥ nreq || throw(ArgumentError("Polynomial order k = $(k) and left,right multiplicities $(ml),$(mr) require a knot sequence of at least $(nreq) points"))
end

"""
    order(t)

Returns the order `k` of the knot set `t`.
"""
order(t::AbstractKnotSet{k}) where k = k

"""
    numinterval(t)

Returns the number of intervals generated by the knot set `t`.

# Examples

```jldoctest
julia> numintervals(LinearKnotSet(3, 0, 1, 2))
2
```
"""
numintervals(t::AbstractKnotSet) = length(t.t)-1

"""
    numfunctions(t)

Returns the number of basis functions generated by knot set `t`.

# Examples

```jldoctest
julia> numfunctions(LinearKnotSet(3, 0, 1, 2))
4

julia> numfunctions(LinearKnotSet(5, 0, 1, 2))
6
```
"""
numfunctions(t::AbstractKnotSet) = length(t) - order(t)

"""
    leftmultiplicity(t)

Return the multiplicity of the left endpoint.
"""
leftmultiplicity(::AbstractKnotSet{k,ml,mr}) where {k,ml,mr} = ml

"""
    rightmultiplicity(t)

Return the multiplicity of the right endpoint.
"""
rightmultiplicity(::AbstractKnotSet{k,ml,mr}) where {k,ml,mr} = mr

"""
    first(t)

Return the first knot of `t`.
"""
first(t::AbstractKnotSet) = first(t.t)

"""
    last(t)

Return the last knot of `t`.
"""
last(t::AbstractKnotSet) = last(t.t)

"""
    length(t)

Return the number of knots of `t`.

# Examples

```jldoctest
julia> length(LinearKnotSet(3, 0, 1, 3))
8

julia> length(LinearKnotSet(3, 0, 1, 3, 1, 1))
4
```
"""
length(t::AbstractKnotSet) = length(t.t) + leftmultiplicity(t) + rightmultiplicity(t) - 2
# size(t) is not very useful for knot sets, but necessary to fulfill
# the interface expected from AbstractVectors.
size(t::AbstractKnotSet) = (length(t),)

"""
    getindex(t, i)

Return the `i`th knot of the knot set `t`, accounting for the
multiplicities of the endpoints.

# Examples

```jldoctest
julia> LinearKnotSet(3, 0, 1, 3)[2]
0.0

julia> LinearKnotSet(3, 0, 1, 3, 1, 1)[2]
0.3333333333333333
```

"""
function getindex(t::AbstractKnotSet{k,ml,mr}, i::Integer) where {k,ml,mr}
    i < 1 && throw(BoundsError(t, i))

    ni = numintervals(t)
    if i < ml
        first(t)
    elseif i < ml + ni
        t.t[i-ml+1]
    elseif i < ml + ni + mr
        last(t)
    else
        throw(BoundsError("Trying to access knot set of length $(length(t)) at index $i"))
    end
end
lastindex(t::AbstractKnotSet) = length(t)
eachindex(t::AbstractKnotSet) = 1:length(t)
# function iterate(t::AbstractKnotSet, (element,i)=(t[1],0))
#     i ≥ length(t) && return nothing
#     element, (t[i+1], i+1)
# end

eltype(t::AbstractKnotSet{k,ml,mr,T}) where {k,ml,mr,T} = T
similar(t::AbstractKnotSet{k,ml,mr,T}) where {k,ml,mr,T} = Vector{T}(undef, length(t))

function show_order_multiplicities(io::IO, t::AbstractKnotSet{k,ml,mr}) where {k,ml,mr}
    write(io, "order k = $(k)")
    k < 7 && write(io, " (",
                  ["constant", "linear", "parabolic", "cubic", "quartic", "quintic"][k],
                  ")")
    ml == k && mr == k && return
    write(io, " (left, right multiplicities: $(ml), $(mr))")
end

function show(io::IO, t::AbstractKnotSet)
    write(io, "$(typeof(t).name)($(eltype(t))) of ")
    show_order_multiplicities(io, t)
    write(io, " on $(first(t)..last(t)) ($(numintervals(t)) intervals)")
end

# TODO: Implement via the `iterate` interface?
"""
    nonempty_intervals(t)

Return the indices of all intervals of the knot set `t` that are
non-empty.

# Examples

```jldoctest
julia> nonempty_intervals(ArbitraryKnotSet(3, [0.0, 1, 1, 3, 4, 6], 1, 3))
4-element Array{Int64,1}:
 1
 3
 4
 5
```

"""
nonempty_intervals(t::AbstractKnotSet) =
    [j for j in 1:length(t)-1
     if t[j] ≠ t[j+1]]

"""
    find_interval(t, x[, i=ml])

Find the interval in the knot set `t` that includes `x`, starting from
interval `i` (which by default is the first non-zero interval of the
knot set). The search complexity is linear, but by storing the result
and using it as starting point for the next call to `find_interval`,
the knot set need only be traversed once.
"""
function find_interval(t::AbstractKnotSet{k,ml,mr,T}, x, i=ml) where {T,k,ml,mr}
    (x < first(t) || x > last(t) || i > length(t) || x < t[i]) && return nothing
    x == last(t) && return length(t) - mr
    for r ∈ i:length(t)
        t[r] > x && return r-1
    end
    @assert false
end

const RightContinuous{T} = Interval{:closed,:open,T}

"""
    within_support(x, t, j)

Return the indices of the elements of `x` that lie withing the compact
support of the `j`th basis function (enumerated `1..n`), given the
knot set `t`. For each index of `x` that is covered, the index `k` of
the interval within which `x[i]` falls is also returned.
"""
function within_support(x::AbstractVector, t::AbstractKnotSet, j::Integer)
    isempty(x) && return 1:0
    k = order(t)
    # The last interval includes the right endpoint as well, whereas
    # the other intervals are only right-continuous.
    IntervalKind(i) = i == numintervals(t) ? ClosedInterval : RightContinuous
    ml = leftmultiplicity(t)
    supports = [(findall(in(IntervalKind(i-ml+1)(t[i], t[i+1])), x), i)
                for i = j:j+k-1]
    filter(s -> !isempty(s[1]), supports)
end

# * Specific knot sets
# ** Arbitrary knot set
"""
    ArbitraryKnotSet{k,ml,mr}(t)

An arbitrary knot set of order `k` and left and right multiplicities
of `ml` and `mr`, respectively. The knot set is specified by the
non-decreasing vector `t`; the same knot may appear multiple times,
which influences the continuity of the B-splines at that location.
"""
struct ArbitraryKnotSet{k,ml,mr,T,V<:AbstractVector{T}} <: AbstractKnotSet{k,ml,mr,T}
    t::V
    ArbitraryKnotSet{k,ml,mr}(t::V) where {k,ml,mr,T,V<:AbstractVector{T}} =
        new{k,ml,mr,T,V}(sort(t))
end

"""
    ArbitraryKnotset(k, t[, ml=k, mr=k])

Construct an order-`k` arbitrary knot set, with the locations of the
knots given by non-decreasing vector `t`.
"""
function ArbitraryKnotSet(k::Integer, t::AbstractVector,
                          ml::Integer=k, mr::Integer=k)
    assert_multiplicities(k, ml, mr, t)
    ArbitraryKnotSet{k,ml,mr}(t)
end

# ** Linear knot set
"""
    LinearKnotSet{k,ml,mr}(t)

A knot set of order `k` and left and right multiplicities of `ml` and
`mr`, respectively, whose knots are uniformly distributed according to
the range `t`. The interior basis functions are thus Cᵏ⁻²-continuous.
"""
struct LinearKnotSet{k,ml,mr,T,R<:AbstractRange{T}} <: AbstractKnotSet{k,ml,mr,T}
    t::R
    LinearKnotSet{k,ml,mr}(t::R) where {k,ml,mr,T,R<:AbstractRange{T}} =
        new{k,ml,mr,T,R}(t)
end

"""
    LinearKnotSet(k, a, b, N[, ml=k, mr=k])

Construct an order-`k` linear knot set spanning from `a` to `b`, with
`N` intervals.

# Examples

```jldoctest
julia> LinearKnotSet(2, 0, 1, 3)
6-element LinearKnotSet{2,2,2,Float64,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}:
 0.0
 0.0
 0.3333333333333333
 0.6666666666666666
 1.0
 1.0
```
"""
function LinearKnotSet(k::Integer, a, b, N::Integer,
                       ml::Integer=k, mr::Integer=k)
    t = range(a, stop=b, length=N+1)
    assert_multiplicities(k, ml, mr, t)
    LinearKnotSet{k,ml,mr}(t)
end

distribution(::LinearKnotSet) = Uniform()

# ** Exponential knot set
"""
    ExpKnotSet{k,ml,mr}(exponents, base, t, include0)

A knot set of order `k` and left and right multiplicities of `ml` and
`mr`, respectively, whose knots are exponentially distributed
according to `t = base .^ exponents`, optionally including 0 as the
left endpoint."""
struct ExpKnotSet{k,ml,mr,T,R<:AbstractRange{T},TV<:AbstractVector{T}} <: AbstractKnotSet{k,ml,mr,T}
    exponents::R
    base::T
    t::TV
    include0::Bool
    ExpKnotSet{k,ml,mr}(exponents::R, base::T, t::TV, include0) where {k,ml,mr,T,R,TV} =
        new{k,ml,mr,T,R,TV}(exponents, base, t, include0)
end

"""
    ExpKnotSet(k, a, b, N[, ml=k, mr=k, base=10, include0=true])

Construct an order-`k` knot spanning from `base^a` to `base^b` in `N`
intervals, optionally including 0 as the left endpoint.

# Examples

```jldoctest
julia> ExpKnotSet(2, -4, 2, 7)
10-element ExpKnotSet{2,2,2,Float64,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}},Array{Float64,1}}:
   0.0
   0.0
   0.0001
   0.001
   0.01
   0.1
   1.0
  10.0
 100.0
 100.0
```
"""
function ExpKnotSet(k::Integer, a::T, b::T, N::Integer,
                    ml::Integer=k, mr::Integer=k;
                    base::T=T(10), include0::Bool=true) where T
    exponents = range(a, stop=b, length=include0 ? N : N+1)
    t = base .^ exponents
    assert_multiplicities(k, ml, mr, t)
    ExpKnotSet{k,ml,mr}(exponents, eltype(t)(base), include0 ? vcat(0,t) : t, include0)
end

function show(io::IO, t::ExpKnotSet)
    write(io, "$(typeof(t).name)($(eltype(t))) of ")
    show_order_multiplicities(io, t)
    write(io, " on ")
    t.include0 && write(io, "0,")
    write(io, "$(t.base^first(t.exponents))..$(t.base^last(t.exponents)) ($(numintervals(t)) intervals)")
end

# function arcsin_knot_set(k::Integer, a::Integer, b::Integer, N::Integer)
#     N2 = N/2
#     arcsin.(range(-N2, stop=N2, lengthN+1)/N2)*(b-a)/π+(a+b)/2
# end

# arcsin_half_knot_set(a, b, N) = arcsin(linspace(0,N,N+1)/N)*(b-a)*2.0/π+a

# function exp_linear_knot_set(a, b, N)
#     ap = a != 0 ? a : 1e-1
#     @assert t[2][1]+t[2][2] == N
#     [logspace(log10(ap),log10(t[1]), t[2][1]); linspace(t[1],b,t[2][2]+1)[2:end]]
# end

@recipe function plot(t::AbstractKnotSet)
    markershape --> :circle
    y = similar(t)
    y′ = 1
    t′ = -Inf
    for (i,e) in enumerate(t)
        y′ = e == t′ ? (y′ + 1) : 1
        t′ = e
        y[i] = y′
    end
    collect(t),y
end

export ArbitraryKnotSet, LinearKnotSet, ExpKnotSet, order, numintervals, numfunctions, nonempty_intervals
