-module(numidx).

-export([add/2, lookup/2, range/2, delete/2]).

-record(node, {
    key = nil,
    value = [],
    lt = nil,
    gt = nil
}).

add(nil, {NewDocId, NewNumidx}) ->
    #node{key=NewNumidx, value=[NewDocId]};
add(#node{key=Numidx, value=Value}=Node, {NewDocId, Numidx}) ->
    Node#node{value=[NewDocId|Value]};
add(#node{key=Numidx}=Node, {_NewDocId, NewNumidx}=NewItem)
        when NewNumidx < Numidx ->
    NewLt = add(Node#node.lt, NewItem),
    Node#node{lt=NewLt};
add(#node{key=Numidx}=Node, {_NewDocId, NewNumidx}=NewItem)
        when NewNumidx > Numidx ->
    NewGt = add(Node#node.gt, NewItem),
    Node#node{gt=NewGt}.


lookup(nil, _SearchKey) ->
    [];
lookup(#node{key=Key, value=Value}, Key) ->
    Value;
lookup(#node{key=Key, lt=Lt}, SearchKey) when SearchKey < Key->
    lookup(Lt, SearchKey);
lookup(#node{key=Key, gt=Gt}, SearchKey) when SearchKey > Key->
    lookup(Gt, SearchKey).


range(nil, _Range) ->
    [];
range(Node, {From, To}) ->
    lists:foldl(fun(Key, Result) ->
       [lookup(Node, Key)|Result]
    end, [], lists:seq(From, To)).


delete(nil, _) ->
    nil;
% Node matches and contains more than one value
delete(#node{key=Numidx, value=Value}=Node, {DelDocId, Numidx})
        when length(Value) > 1 ->
    Node#node{value=lists:delete(DelDocId, Value)};
% Node matches and has no child nodes
delete(#node{key=Numidx, lt=nil, gt=nil}, {_DelDocId, Numidx}) ->
    nil;
% Node matches and has one child (lt)
delete(#node{key=Numidx, lt=Lt, gt=nil}, {_DelDocId, Numidx}) ->
    Lt;
% Node matches and has one child (gt)
delete(#node{key=Numidx, lt=nil, gt=Gt}, {_DelDocId, Numidx}) ->
    Gt;
% Node matches and has both child nodes
delete(#node{key=Numidx, lt=Lt, gt=Gt}, {_DelDocId, Numidx}) ->
    % Find greatest node an rebuild tree (thanks Concurrent programming in
    % Erlang Part I pp. 61)
    NewNode = deletesp(Lt),
    NewNode#node{gt=Gt};
% Node doesn't match, move on
delete(#node{key=Numidx}=Node, {_DelDocId, DelNumidx}=DelItem)
        when DelNumidx < Numidx ->
    DelLt = delete(Node#node.lt, DelItem),
    Node#node{lt=DelLt};
% Node doesn't match, move on
delete(#node{key=Numidx}=Node, {_DelDocId, DelNumidx}=DelItem)
        when DelNumidx > Numidx ->
    DelGt = delete(Node#node.gt, DelItem),
    Node#node{gt=DelGt}.

% Abusing #node, always write the result into #node.lt
deletesp(#node{gt=nil}=Node) ->
    Node;
deletesp(#node{gt=Gt}=Node) ->
    NewNode = deletesp(Gt),
    NewNode#node{lt=Node#node{gt=NewNode#node.lt}}.