%%% @doc
%%% kvex — pure Erlang approximate k-NN vector search on the BEAM.
%%%
%%% Two ETS tables per index:
%%%   • vec table  — `{Id, F32Bin, BinVec}' per vector (source of truth)
%%%   • flat cache — `{flat, F32FlatBin, BvecFlatBin, IdsTuple}' rebuilt on
%%%                  every insert; a single refc-binary per flat, so search
%%%                  never iterates over individual Erlang terms.
%%%
%%% Search path (two sied NIF calls on flat binaries, no Erlang list work):
%%%   1. `sied:hamming_topk_flat/4'   — SIMD POPCNT on BvecFlat, O(N)+O(K logK)
%%%   2. `sied:dot_product_topk_flat/4' — SIMD dot-product on F32Flat candidates
%%%
%%% == Quick start ==
%%% ```
%%% {ok, Ix} = kvex:new(128),
%%% Vec      = [rand:uniform() || _ <- lists:seq(1, 128)],
%%% ok       = kvex:add(Ix, 42, Vec),
%%% {ok, Rs} = kvex:search(Ix, Vec, 5),
%%% ok       = kvex:delete(Ix).
%%% '''
%%% @end
-module(kvex).

-export([version/0, new/1, new/2, delete/1, size/1,
         add/3, add_batch/2, search/3,
         normalize/1, cosine_search/3]).

-define(OVERSAMPLE, 10).

%% flat cache key stored in the vec table (atom, never clashes with id())
-define(FLAT_KEY, '$kvex_flat').

-opaque index() :: #{table := ets:tid(), dim := pos_integer()}.
-type opts()    :: #{bits => 2 | 3 | 4}.
-type id()      :: non_neg_integer() | binary().
-type vector()  :: [float()] | binary().

-export_type([index/0, opts/0, id/0, vector/0]).

%%%===================================================================
%%% Public API
%%%===================================================================

-spec version() -> binary().
version() -> <<"0.2.0">>.

-spec new(Dim :: pos_integer()) -> {ok, index()} | {error, term()}.
new(Dim) -> new(Dim, #{}).

-spec new(Dim :: pos_integer(), opts()) -> {ok, index()} | {error, term()}.
%% @doc Creates an empty index for vectors of dimension `Dim'.
new(Dim, _Opts) when is_integer(Dim), Dim > 0 ->
    Tid = ets:new(kvex, [set, protected]),
    ets:insert(Tid, {?FLAT_KEY, <<>>, <<>>, {}}),
    {ok, #{table => Tid, dim => Dim}};
new(Dim, _Opts) ->
    {error, {bad_dim, Dim}}.

-spec delete(index()) -> ok.
delete(#{table := Tid}) ->
    ets:delete(Tid),
    ok.

-spec size(index()) -> non_neg_integer().
size(#{table := Tid}) ->
    ets:info(Tid, size) - 1.   % subtract the flat_cache sentinel

-spec add(index(), id(), vector()) -> ok | {error, term()}.
%% @doc Inserts a single vector. Rebuilds the flat cache — O(N) copy.
add(#{table := Tid, dim := Dim}, Id, Vec0) ->
    F32Bin = to_f32_bin(Vec0),
    case byte_size(F32Bin) div 4 of
        Dim ->
            {ok, BinVec} = sied:to_binary_f32_bin(F32Bin),
            ets:insert(Tid, {Id, F32Bin, BinVec}),
            rebuild_flat(Tid),
            ok;
        Got ->
            {error, {dim_mismatch, Dim, Got}}
    end.

-spec add_batch(index(), [{id(), vector()}]) -> ok | {error, term()}.
%% @doc Inserts vectors in batch. Builds the flat binary incrementally — O(batch).
add_batch(#{table := Tid, dim := Dim}, Pairs) when is_list(Pairs) ->
    case build_entries(Pairs, Dim, 0, [], [], [], []) of
        {ok, Entries, F32New, BvecNew, IdsNew} ->
            ets:insert(Tid, Entries),
            [{?FLAT_KEY, F32Old, BvecOld, IdsOldT}] = ets:lookup(Tid, ?FLAT_KEY),
            IdsOld   = tuple_to_list(IdsOldT),
            F32Flat  = <<F32Old/binary,  F32New/binary>>,
            BvecFlat = <<BvecOld/binary, BvecNew/binary>>,
            IdsTuple = list_to_tuple(IdsOld ++ IdsNew),
            ets:insert(Tid, {?FLAT_KEY, F32Flat, BvecFlat, IdsTuple}),
            ok;
        {error, _} = Err ->
            Err
    end.

-spec search(index(), Query :: vector(), K :: pos_integer()) ->
        {ok, [{id(), Score :: float()}]} | {error, term()}.
search(#{table := Tid, dim := Dim}, Query0, K)
        when is_integer(K), K > 0 ->
    QBin = to_f32_bin(Query0),
    case byte_size(QBin) div 4 of
        Dim ->
            [{?FLAT_KEY, F32Flat, BvecFlat, IdsTuple}] = ets:lookup(Tid, ?FLAT_KEY),
            N = tuple_size(IdsTuple),
            case N of
                0 -> {error, empty_index};
                _ -> do_search(QBin, F32Flat, BvecFlat, IdsTuple, K, N)
            end;
        Got ->
            {error, {dim_mismatch, Dim, Got}}
    end.

-spec normalize(vector()) -> {ok, [float()]} | {error, term()}.
normalize(Vec) when is_list(Vec)   -> sied:l2_normalize_f32(Vec);
normalize(Vec) when is_binary(Vec) -> sied:l2_normalize_f32(f32_bin_to_list(Vec)).

-spec cosine_search(index(), Query :: vector(), K :: pos_integer()) ->
        {ok, [{id(), Score :: float()}]} | {error, term()}.
cosine_search(Ix, Query, K) when is_list(Query), is_integer(K), K > 0 ->
    case sied:l2_normalize_f32(Query) of
        {ok, NormQ} -> search(Ix, NormQ, K);
        Error       -> Error
    end;
cosine_search(Ix, Query, K) when is_binary(Query), is_integer(K), K > 0 ->
    cosine_search(Ix, f32_bin_to_list(Query), K).

%%%===================================================================
%%% Internal
%%%===================================================================

do_search(QBin, F32Flat, BvecFlat, IdsTuple, K, N) ->
    {ok, QQuantBin} = sied:to_binary_f32_bin(QBin),
    VecBLen   = byte_size(QQuantBin),
    VecF32Len = byte_size(QBin),
    CandCount = min(K * ?OVERSAMPLE, N),
    %% Phase 1 — SIMD POPCNT on flat binary, returns top-CandCount indices
    {ok, CandIdxs} = sied:hamming_topk_flat(QQuantBin, BvecFlat, VecBLen, CandCount),
    %% Phase 2 — SIMD dot-product on flat f32, returns [{Score, Idx}] sorted desc
    {ok, Scored} = sied:dot_product_topk_flat(QBin, F32Flat, VecF32Len, CandIdxs),
    {ok, [{element(Idx + 1, IdsTuple), Score}
          || {Score, Idx} <- lists:sublist(Scored, K)]}.

%% Rebuild flat cache from all records in the vec table.
%% Called after single add/3 — O(N) scan + binary concat.
rebuild_flat(Tid) ->
    All = ets:select(Tid, [{{'$1','$2','$3'}, [{'/=','$1',{const,?FLAT_KEY}}], [{{'$1','$2','$3'}}]}]),
    {F32Flat, BvecFlat, Ids} = lists:foldl(
        fun({Id, F32, BV}, {F, B, Is}) ->
            {<<F/binary, F32/binary>>, <<B/binary, BV/binary>>, [Id | Is]}
        end,
        {<<>>, <<>>, []},
        All
    ),
    ets:insert(Tid, {?FLAT_KEY, F32Flat, BvecFlat, list_to_tuple(lists:reverse(Ids))}).

build_entries([], _Dim, _Pos, EAcc, F32Acc, BvAcc, IdsAcc) ->
    {ok, lists:reverse(EAcc),
     iolist_to_binary(lists:reverse(F32Acc)),
     iolist_to_binary(lists:reverse(BvAcc)),
     lists:reverse(IdsAcc)};
build_entries([{Id, Vec0} | Rest], Dim, Pos, EAcc, F32Acc, BvAcc, IdsAcc) ->
    F32Bin = to_f32_bin(Vec0),
    case byte_size(F32Bin) div 4 of
        Dim ->
            {ok, BinVec} = sied:to_binary_f32_bin(F32Bin),
            build_entries(Rest, Dim, Pos + 1,
                [{Id, F32Bin, BinVec} | EAcc],
                [F32Bin  | F32Acc],
                [BinVec  | BvAcc],
                [Id      | IdsAcc]);
        Got ->
            {error, {dim_mismatch, Pos, Dim, Got}}
    end.

to_f32_bin(Vec) when is_binary(Vec) -> Vec;
to_f32_bin(Vec) when is_list(Vec)   ->
    << <<F:32/float-little>> || F <- Vec >>.

f32_bin_to_list(<<>>) -> [];
f32_bin_to_list(<<F:32/float-little, Rest/binary>>) ->
    [F | f32_bin_to_list(Rest)].