import Base: length, getindex
export SampleSet, MetricNN, QuasiMetricNN
export NearNeighborCache, MutableNNC, ImmutableNNC
export ControlCache, EmptyControlCache, MutableControlCache, ImmutableControlCache
export DistanceDataStructure, EmptyDistanceDS, BruteDistanceDS, TreeDistanceDS
export filter_neighborhood, addpoints
# export resetcache!
# export saveNN, loadNN!
export inball, knn, inball!, knn!, mutualknn!         # for metrics only (symmetric)
export inballF, knnF, inballF!, knnF!, mutualknnF!    # forwards
export inballB, knnB, inballB!, knnB!, mutualknnB!    # backwards

### Typedefs
abstract SampleSet
abstract NearNeighborCache{T}; begin
    ## NearNeighborCache implementations
    # TODO: consider EmptyNNC for online algorithms or streaming-batch
    immutable MutableNNC{T} <: NearNeighborCache{T}
        D::Vector{SparseVector{T,Int}}
        r::Vector{T}
    end
    MutableNNC{T}(N::Int, ::Type{T}) = MutableNNC{T}(Array(SparseVector{T,Int}, N), zeros(T, N))
    immutable ImmutableNNC{T} <: NearNeighborCache{T}
        D::SparseMatrixCSC{T,Int}
        r::Vector{T}
    end
    ImmutableNNC(M::MutableNNC) = ImmutableNNC(hcat(M.D...), M.r)
    MutableNNC{T}(I::ImmutableNNC{T}) = MutableNNC(size(I.D,1), T) # clear the cache
end

abstract ControlCache{U<:ControlInfo}; begin
    ## ControlCache implementations 
    immutable EmptyControlCache{U} <: ControlCache{U} end
    EmptyControlCache{U}(::Type{U}) = EmptyControlCache{U}()
    EmptyControlCache() = EmptyControlCache(NullControl)
    immutable MutableControlCache{U} <: ControlCache{U}
        US::Vector{SparseVector{U,Int}}
    end
    MutableControlCache{U}(N::Int, ::Type{U}) = MutableControlCache{U}(Array(SparseVector{U,Int}, N))
    immutable ImmutableControlCache{U} <: ControlCache{U}
        US::SparseMatrixCSC{U,Int}
    end
    ImmutableControlCache(M::MutableControlCache) = ImmutableControlCache(hcat(M.U...))
    MutableControlCache{U}(I::ImmutableControlCache{U}) = MutableControlCache(size(I.US,1), U) # clear the cache
end

abstract DistanceDataStructure{T}; begin
    ## DistanceDataStructure implementations
    immutable EmptyDistanceDS{T} <: DistanceDataStructure{T} end
    EmptyDistanceDS{T}(::Type{T}) = EmptyDistanceDS{T}()
    immutable BruteDistanceDS{A<:AbstractMatrix,T} <: DistanceDataStructure{T}
        Dmat::A
    end
    BruteDistanceDS{T}(Dmat::AbstractMatrix{T}) = BruteDistanceDS{typeof(Dmat),T}(Dmat)
    immutable TreeDistanceDS{NNT<:NearestNeighbors.NNTree,T} <: DistanceDataStructure{T}
        tree::NNT
    end
    TreeDistanceDS{V}(tree::NearestNeighbors.NNTree{V}) = TreeDistanceDS{typeof(tree),eltype(V)}(tree)
end

### Construction
type MetricNN{M<:SymmetricDistance,S<:State,N<:NearNeighborCache,D<:DistanceDataStructure,U<:ControlCache} <: SampleSet
    V::Vector{S}
    dist::M
    init::S
    cache::N
    DS::D
    US::U
end
function MetricNN{S<:State}(V::Vector{S}, dist::SymmetricDistance, init::S)
    N = length(V)
    T = eltype(S)
    MetricNN(V, dist, init, MutableNNC(N,T), helper_data_structures(V, dist)...)
end

type QuasiMetricNN{M<:AsymmetricDistance,S<:State,N<:NearNeighborCache,
                   D<:DistanceDataStructure,U<:ControlCache} <: SampleSet
    V::Vector{S}
    dist::M
    init::S
    cacheF::N   # Everything's just pointers, but F and B may be different (e.g. CSC and CSR), trading memory for speed
    cacheB::N
    DSF::D
    USF::U
    DSB::D
    USB::U
end
function QuasiMetricNN{S<:State}(V::Vector{S}, dist::AsymmetricDistance, init::S)
    N = length(V)
    T = eltype(S)
    QuasiMetricNN(V, dist, init, MutableNNC(N,T), MutableNNC(N,T), helper_data_structures(V, dist)...)
end

## Fallbacks (should usually be extended)
helper_data_structures{S}(V::Vector{S}, d::SymmetricDistance) = (EmptyDistanceDS(eltype(S)),
                                                                 EmptyControlCache(controltype(d)))
function helper_data_structures{S}(V::Vector{S}, d::AsymmetricDistance)
    DS, US = EmptyDistanceDS(eltype(S)), EmptyControlCache(controltype(d))
    (DS, US, DS, US)
end

### Utilities

@inline function filter_neighborhood(n::AbstractSparseVector, f::BitVector)
    inds, ds = nonzeroinds(n), nonzeros(n)
    SparseVector(length(n), inds[f[inds]], ds[f[inds]])
end
addpoints(NN::MetricNN, W) = MetricNN([NN.V; W], NN.dist, NN.init)
addpoints(NN::QuasiMetricNN, W) = QuasiMetricNN([NN.V; W], NN.dist, NN.init)
length(NN::SampleSet) = length(NN.V)
getindex(NN::SampleSet, i::Int) = i > 0 ? NN.V[i] : NN.init
getindex(NN::SampleSet, I::AbstractVector) = [NN[i] for i in I]

### Serialization



### Near Neighbor Computation
# TODO: forwards/backwards should be probably moved into the datastructure; perhaps a type parameterization
@inline inball!(NN::MetricNN, v::Int, r, f::BitVector) = filter_neighborhood(inball!(NN, v, r), f)
@inline inballF!(NN::QuasiMetricNN, v::Int, r, f::BitVector) = filter_neighborhood(inballF!(NN, v, r), f)
@inline inballB!(NN::QuasiMetricNN, v::Int, r, f::BitVector) = filter_neighborhood(inballB!(NN, v, r), f)
for direction in ("", "F", "B")
    NNsym = (direction == "" ? :MetricNN : :QuasiMetricNN)
    inballD! = Symbol(:inball, direction, "!")
    cacheD = Symbol(:cache, direction)
    DSD = Symbol(:DS, direction)
    @eval $inballD!{M,S,N<:ImmutableNNC}(NN::$NNsym{M,S,N}, v::Int, r) = viewcol(NN.$cacheD.D, v)
    @eval function $inballD!{M,S,N<:MutableNNC}(NN::$NNsym{M,S,N}, v::Int, r)
        if !isdefined(NN.$cacheD.D, v)
            NN.$cacheD.D[v] = inball(NN.V, NN.dist, NN.$DSD, v, r, $(direction != "B"))
            NN.$cacheD.r[v] = r
        end
        NN.$cacheD.D[v]
    end
end

function inball(V::Vector, dist::PreMetric, DS::DistanceDataStructure, v::Int, r, forwards::Bool = true)
    allds = forwards ? colwise(dist, V[v], V) : colwise(dist, V, V[v])
    N = length(V)
    inds = Vector{Int}(0)
    ds = similar(allds, 0)
    @inbounds for i in 1:N
        if i != v && allds[i] <= r
            push!(inds, i)
            push!(ds, allds[i])
        end
    end
    SparseVector(N, inds, ds)
end

function inball{A<:Matrix}(V::Vector, dist::PreMetric, DS::BruteDistanceDS{A}, v::Int, r, forwards::Bool = true)
    N = length(V)
    inds = Vector{Int}(0)
    ds = similar(DS.Dmat, 0)
    @inbounds for i in 1:N
        if i != v && DS.Dmat[i,v] <= r
            push!(inds, i)
            push!(ds, allds[i])
        end
    end
    SparseVector(N, inds, ds)
end

function inball{A<:SparseMatrixCSC}(V::Vector, dist::PreMetric, DS::BruteDistanceDS{A}, v::Int, r, forwards::Bool = true)
    vcol = viewcol(DS.Dmat, v)
    allinds, allds = nonzeroinds(vcol), nonzeros(vcol)
    inds = Vector{Int}(0)
    ds = Vector{eltype(A)}(0)
    @inbounds for i in 1:length(allinds)
        if allinds[i] != v && allds[i] <= r
            push!(inds, allinds[i])
            push!(ds, allds[i])
        end
    end
    SparseVector(length(V), inds, ds)
end

function inball{S}(V::Vector{S}, dist::PreMetric, DS::TreeDistanceDS, v::Int, r, forwards::Bool = true)
    inds = inrange(DS.tree, V[v], r, true)
    inds = deleteat!(inds, searchsortedfirst(inds, v))
    SparseVector(length(V), inds, forwards ? colwise(dist, V[v], V[inds]) : colwise(dist, V[inds], V[v]))
end

function inball{S}(V::Vector{S}, dist::ChoppedPreMetric, DS::TreeDistanceDS, v::Int, r, forwards::Bool = true)
    allinds = inrange(DS.tree, V[v], r, true)
    allinds = deleteat!(allinds, searchsortedfirst(allinds, v))
    allds = forwards ? colwise(dist, V[v], V[allinds]) : colwise(dist, V[allinds], V[v])
    inds = Vector{Int}(0)
    ds = similar(allds, 0)
    @inbounds for i in 1:length(allinds)
        if allds[i] <= r
            push!(inds, allinds[i])
            push!(ds, allds[i])
        end
    end
    SparseVector(length(V), inds, ds)
end

for ff in [x -> Symbol("inball", x), x -> Symbol("inball", x, "!")]
    @eval $(ff("F"))(NN::MetricNN, args...) = $(ff(""))(NN, args...)
    @eval $(ff("B"))(NN::MetricNN, args...) = $(ff(""))(NN, args...)
end
